astro-nest 0.7.2.dev2__tar.gz → 0.7.2.dev4__tar.gz

{astro_nest-0.7.2.dev2 → astro_nest-0.7.2.dev4}/NEST/core.py

@@ -14,6 +14,8 @@ except ImportError:
     has_requests = False
 try:
     import matplotlib.pyplot as plt
+    from matplotlib.patches import Polygon
+    from matplotlib.colors import ListedColormap
     has_matplotlib = True
 except ImportError:
     has_matplotlib = False
@@ -33,6 +35,16 @@ def custom_warning(message, category, filename, lineno, file=None, line=None):
 
 warnings.showwarning = custom_warning
 
+def import_cmaps():
+    cmaps = {}
+    if has_matplotlib:
+        cmaps_colors = dict(np.load(os.path.join(NEST_DIR, 'cmaps.npz')))
+        for name, colors in cmaps_colors.items():
+            cmaps[name] = ListedColormap(colors)
+    return cmaps
+
+cmaps = import_cmaps()
+
 def sanitize_input(input):
     if input is not None and type(input) is not list:
         if hasattr(input,'ndim') and input.ndim == 0:
@@ -153,6 +165,8 @@ class AgeModel:
         self.means = None
         self.modes = None
         self.stds = None
+        self.pop_age = None
+        self.pop_age_std = None
         self.load_neural_network(self.model_name)
 
     def __repr__(self):
@@ -391,7 +405,7 @@ class AgeModel:
         
         return in_domain
     
-    def population_age(self,nbins=280,min_age=0,max_age=14,check_domain=True,n_mc=100):
+    def population_age(self,nbins=280,min_age=0,max_age=14,check_domain=True,n_mc=100,use_tqdm=True):
         if self.ages is None:
             raise ValueError('No samples have been stored yet. Make sure to run ages_prediction() with store_samples=True')
         
@@ -407,15 +421,19 @@ class AgeModel:
             ages = self.ages
 
         ages = np.array([age for age in ages if np.any(np.isnan(age)) == False])
+
+        if len(ages) == 0:
+            raise ValueError('No valid age samples available to compute population age (maybe all stars are out of domain ?)')
         
         if ages.shape[1] == 1:
             return np.median(ages)
 
         pdfs = []
+        best_ages = []
         loop = range(n_mc)
-        if self.use_tqdm:
+        if self.use_tqdm and use_tqdm:
             loop = tqdm(loop)
-        min_offset = 0.1
+        epsilon = np.clip(1/ages.shape[0],0.001,0.1)
         for i in loop:
             random_star_indices = np.random.choice(list(range(ages.shape[0])), size=ages.shape[0], replace=True)
             random_star_indices = random_star_indices.astype(int)
@@ -425,15 +443,24 @@ class AgeModel:
                 age = random_star_ages[j]
                 pdf, bins = np.histogram(age,bins=nbins,range=(min_age,max_age))
                 pdf = pdf.astype(float)
-                pdf += min_offset
+                pdf += epsilon
                 global_pdf *= pdf
             best_age = bins[np.argmax(global_pdf)] + (bins[1]-bins[0])/2
             if best_age != bins[0]+(bins[1]-bins[0])/2:
+                best_ages.append(best_age)
                 pdfs.append(bins[np.argmax(global_pdf)] + (bins[1]-bins[0])/2)
 
+        if len(pdfs) == 0:
+            warnings.warn('No valid population age could be computed from the samples')
+            return np.nan
+
         population_age = np.median(pdfs)
+        population_age_std = np.std(best_ages)
 
-        return population_age
+        self.pop_age = population_age
+        self.pop_age_std = population_age_std
+
+        return population_age, population_age_std
     
     def propagate(self,X,neural_network,scaler):
         if self.use_sklearn:
@@ -499,7 +526,14 @@ class AgeModel:
         self.stds = np.std(self.ages,axis=1)
         return self.stds
     
-    def HR_diagram(self,isochrone_met=0,plot_isochrone=True,plot_stars=True,age_type='median',check_domain=True,fig=None,ax=None,**kwargs):
+    def HR_diagram(self,
+                   isochrone_met=0,
+                   plot_isochrone=True,plot_stars=True,
+                   isochrone_ages=None,isochrone_ages_std=None,
+                   age_type='median',
+                   check_domain=True,
+                   fig=None,ax=None,
+                   **kwargs):
         if has_matplotlib == False:
             raise ImportError('matplotlib is required for HR diagram plotting')
         
@@ -584,6 +618,8 @@ class AgeModel:
             kwargs['color'] = 'k'
         if 'lw' not in kwargs and 'linewidth' not in kwargs:
             kwargs['linewidth'] = 0.5
+        if 'alpha' not in kwargs:
+            kwargs['alpha'] = 0.25
         for i,isochrone in enumerate(isochrones):
             iso_age = isochrone['age']
             iso_mag = isochrone['MG']
@@ -599,7 +635,6 @@ class AgeModel:
             ages.append(iso_age)
             colors.append('k')
 
-        isochrone_ages = []
         vmin = None
         vmax = None
         cb = None
@@ -609,37 +644,39 @@ class AgeModel:
                 col,
                 mag,
                 c=age,
-                cmap='viridis',
+                cmap=cmaps['tempo_R'],
                 zorder=4,
                 marker='*'
             )
             ages_arr = np.array(age)
             data_vmin = np.nanmin(ages_arr)
             data_vmax = np.nanmax(ages_arr)
-            # If any values fall outside [0,14], restrict the colorbar to [0,14]
+
             if data_vmin < 0 or data_vmax > 14:
                 vmin, vmax = 0.0, 14.0
             else:
                 vmin, vmax = float(data_vmin), float(data_vmax)
-            # Ensure a valid range
+
             if vmin >= vmax:
                 vmin, vmax = max(0.0, vmin - 0.5), min(14.0, vmax + 0.5)
                 if vmin >= vmax:
                     vmin, vmax = 0.0, 14.0
             scatter.set_clim(vmin, vmax)
             cb = plt.colorbar(scatter, label='Age [Gyr]')
-            if (type(plot_isochrone) is bool and plot_isochrone != False) or type(plot_isochrone) in (float,int,list):
-                #if self.model_name != 'BaSTI':
-                    #raise ValueError('Plotting individual isochrones is only available for the BaSTI model at the moment')
-                if type(plot_isochrone) is bool and plot_isochrone:
-                    if self.ages is not None and self.ages.shape[1] > 1:
-                        isochrone_ages = [self.population_age()]
-                    else:
-                        isochrone_ages = [np.median(age)]
-        if type(plot_isochrone) is float or type(plot_isochrone) is int:
-            isochrone_ages = [float(plot_isochrone)]
-        elif type(plot_isochrone) is list:
-            isochrone_ages = plot_isochrone
+        
+        if plot_isochrone and isochrone_ages is None and isochrone_ages_std is None:
+            if self.ages is not None and self.ages.shape[1] > 1:
+                if self.pop_age is None:
+                    self.population_age(check_domain=check_domain,use_tqdm=False)
+                isochrone_ages = [self.pop_age]
+                isochrone_ages_std = [self.pop_age_std]
+            else:
+                isochrone_ages = [np.median(age)]
+                isochrone_ages_std = [np.std(age)]
+        if type(isochrone_ages) is float or type(isochrone_ages) is int:
+            isochrone_ages = [float(isochrone_ages)]
+        if type(isochrone_ages_std) is float or type(isochrone_ages_std) is int:
+            isochrone_ages_std = [float(isochrone_ages_std)]
         
         if len(isochrone_ages) > 0:
             if vmin is None or vmax is None:
@@ -649,30 +686,37 @@ class AgeModel:
                 else:
                     vmin = max(0, isochrone_ages[0] - 0.5)
                     vmax = min(14, isochrone_ages[0] + 0.5)
-            for isochrone_age in isochrone_ages:
+            for i,isochrone_age in enumerate(isochrone_ages):
                 isochrone = self.get_closest_isochrone(isochrone_age)
-                norm = plt.Normalize(vmin, vmax)
-                cmap = plt.get_cmap('viridis')
-                color = cmap(norm(isochrone_age))
+                color = cmaps['acton']((i+1)/(len(isochrone_ages)+1))
+
+                if isochrone_ages_std is not None and len(isochrone_ages_std) == len(isochrone_ages):
+                    isochrone_age_std = isochrone_ages_std[isochrone_ages.index(isochrone_age)]
+                    label = r'$\tau$' + f'={isochrone_age:.2f} ' + r'$\pm$' + f' {isochrone_age_std:.2f} Gyr'
+                    isochrone_low = self.get_closest_isochrone(isochrone_age - isochrone_age_std)
+                    isochrone_upp = self.get_closest_isochrone(isochrone_age + isochrone_age_std)
+                    verts = np.concatenate([
+                        np.column_stack([isochrone_low['BP-RP'], isochrone_low['MG']]),
+                        np.column_stack([isochrone_upp['BP-RP'][::-1], isochrone_upp['MG'][::-1]])
+                    ])
+                    poly = Polygon(verts, closed=True, facecolor=color, alpha=0.25, edgecolor=None, zorder=6)
+                    ax.add_patch(poly)
+                else:
+                    label = r'$\tau$' + f'={isochrone_age:.2f} Gyr'
+
                 line, = ax.plot(
                     isochrone['BP-RP'],
                     isochrone['MG'],
                     color=color,
                     linewidth=2,
                     zorder=5,
-                    label=f'Isochrone (age={isochrone_age:.2f} Gyr)'
+                    label=label
                 )
+
                 has_labels = True
                 lines.append(line)
                 ages.append(isochrone_age)
                 colors.append(color)
-            
-            if cb is None:
-                norm = plt.Normalize(vmin, vmax)
-                mappable = plt.cm.ScalarMappable(norm=norm, cmap='viridis')
-                mappable.set_array([vmin,vmax])
-                cb = plt.colorbar(mappable, ax=ax, label='Age [Gyr]')
-                has_labels = True
 
         annot = ax.annotate("", xy=(0,0), xytext=(0,15), textcoords="offset points",
                             ha='center',

{astro_nest-0.7.2.dev2 → astro_nest-0.7.2.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: astro-nest
-Version: 0.7.2.dev2
+Version: 0.7.2.dev4
 Summary: A python package that provides neural networks to estimate ages of stars.
 Author-email: Tristan Boin <tristan.boin@obspm.fr>
 License-Expression: MIT

{astro_nest-0.7.2.dev2 → astro_nest-0.7.2.dev4}/astro_nest.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: astro-nest
-Version: 0.7.2.dev2
+Version: 0.7.2.dev4
 Summary: A python package that provides neural networks to estimate ages of stars.
 Author-email: Tristan Boin <tristan.boin@obspm.fr>
 License-Expression: MIT

{astro_nest-0.7.2.dev2 → astro_nest-0.7.2.dev4}/astro_nest.egg-info/SOURCES.txt

@@ -2,6 +2,7 @@ pyproject.toml
 readme.md
 NEST/__init__.py
 NEST/_version.py
+NEST/cmaps.npz
 NEST/core.py
 NEST/domain.pkl
 NEST/tutorial.ipynb

{astro_nest-0.7.2.dev2 → astro_nest-0.7.2.dev4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "astro-nest"
-version = "0.7.2dev2"
+version = "0.7.2dev4"
 authors = [
   { name="Tristan Boin", email="tristan.boin@obspm.fr" },
 ]
@@ -27,4 +27,4 @@ include = ["NEST"]
 Homepage = "https://star-age.github.io/NEST-docs/"
 
 [tool.setuptools.package-data]
-NEST = ["NEST.py","models/**/*","domain.pkl","tutorial.ipynb"]
+NEST = ["NEST.py","models/**/*","domain.pkl","tutorial.ipynb","cmaps.npz"]