finally getting around to fixing the retrieval of the stellar atmosph…

…ere spectra
Ciela-Institute · Dec 7, 2023 · d7eac3f · d7eac3f
1 parent d3a6be2
commit d7eac3f
Show file tree

Hide file tree

Showing 2 changed files with 108 additions and 17 deletions.
diff --git a/lighthouse/stellar_atmosphere_spectrum/get_stellar_templates.py b/lighthouse/stellar_atmosphere_spectrum/get_stellar_templates.py
@@ -49,11 +49,17 @@ def get_polynomial_coefficients_villaume2017a():
 
 
     B = {
-        "Cool_Dwarfs": {"surface_gravity": (4.0, 6.0),  "effective_temperature": (2500,4000)},
-        "Cool_Giants": {"surface_gravity": (-0.5, 4.0), "effective_temperature": (2500,4000)},
-        "Warm_Dwarfs": {"surface_gravity": (4.0, 6.0),  "effective_temperature": (4000,6000)},
-        "Warm_Giants": {"surface_gravity": (-0.5, 4.0), "effective_temperature": (4000,6000)},
-        "Hot_Stars":   {"surface_gravity": (-0.5, 6),   "effective_temperature": (6000,12000)},
+        "Cool_Giants":    {"surface_gravity": (-0.5, 4.0), "effective_temperature": (2500,4000)},
+        "Warm_Giants":    {"surface_gravity": (-0.5, 4.0), "effective_temperature": (4500,5500)},
+        "Hottish_Giants": {"surface_gravity": (-0.5, 4.0), "effective_temperature": (5500,6500)},
+        "Coolish_Giants": {"surface_gravity": (-0.5, 4.0), "effective_temperature": (3500,4500)},
+
+        "Cool_Dwarfs":    {"surface_gravity": (4.0, 10.0),  "effective_temperature": (2500,3000)},
+        "Warm_Dwarfs":    {"surface_gravity": (4.0, 10.0),  "effective_temperature": (5500,6000)},
+        "Coolish_Dwarfs": {"surface_gravity": (4.0, 10.0),  "effective_temperature": (3000,5500)},
+
+        "Hot_Giants":     {"surface_gravity": (-0.5, 4.0),   "effective_temperature": (6500,12000)},
+        "Hot_Dwarfs":     {"surface_gravity": (4.0, 10.0),   "effective_temperature": (6000,12000)},
     }
 
     with open(os.path.join(data_path, 'bounds.dat'), "w") as f:

diff --git a/lighthouse/stellar_atmosphere_spectrum/polynomial_evaluator.py b/lighthouse/stellar_atmosphere_spectrum/polynomial_evaluator.py
@@ -38,12 +38,30 @@ def __init__(self):
                 self.wavelength         = torch.tensor(coeffs.to_numpy()[:,0], dtype  = torch.float64)
                 self.reference[name]    = torch.tensor(coeffs.to_numpy()[:,1], dtype  = torch.float64)
                 self.coefficients[name] = torch.tensor(coeffs.to_numpy()[:,2:], dtype = torch.float64)
-                
+
     def get_spectrum(self, teff, logg, feh) -> torch.Tensor:
+
         """
-        Setting up some boundaries
+        These weights are used later to ensure
+        smooth behavior.
         """
+        # Overlap of cool dwarf and warm dwarf training sets
+        d_teff_overlap = np.linspace(3000, 5500, num=100)
+        d_weights = np.linspace(1, 0, num=100)
+
+        # Overlap of warm giant and hot star training sets
+        gh_teff_overlap = np.linspace(5500, 6500, num=100)
+        gh_weights = np.linspace(1, 0, num=100)
+
+        # Overlap of warm giant and cool giant training sets
+        gc_teff_overlap = np.linspace(3500, 4500, num=100)
+        gc_weights = np.linspace(1, 0, num=100)
 
+
+
+        """
+        Setting up some boundaries
+        """
         teff2 = teff
         logg2 = logg
         if teff2 <= 2800.:
@@ -52,26 +70,93 @@ def get_spectrum(self, teff, logg, feh) -> torch.Tensor:
             logg2 = torch.tensor(-0.5)
 
         # Normalizing to solar values
-        # logt = np.log10(teff2) - 3.7617
-        logt = np.log10(teff) - 3.7617
+        logt = np.log10(teff2) - 3.7617
+        #logt = np.log10(teff) - 3.7617
         logg = logg - 4.44
 
-        # print(logg2, teff2)
         for key, ranges in self.bounds.items():
             if ranges["surface_gravity"][0] <= logg2 <= ranges["surface_gravity"][1] and ranges["effective_temperature"][0] <= teff2 <= ranges["effective_temperature"][1]:
-                stellar_type = key
+
+                if key is 'Hot_Giants' or key is 'Hot_Dwarfs':
+                    stellar_type = 'Hot_Stars'
+                else:
+                    stellar_type = key
                 break
-        else:
-            stellar_type = "Cool_Giants"
+
 
         K  = torch.stack((torch.as_tensor(logt, dtype = torch.float64),
                           torch.as_tensor(feh,  dtype = torch.float64),
                           torch.as_tensor(logg, dtype = torch.float64)))
-        PP = torch.as_tensor(self.polynomial_powers[stellar_type], dtype = torch.float64)
-        X  = torch.prod(K**PP, dim = -1)
 
-        flux = np.exp(self.coefficients[stellar_type] @ X)
-        flux *= self.reference[stellar_type]
+        easy_types = ['Cool_Giants', 'Warm_Giants', 'Cool_Dwarfs', 'Warm_Dwarfs', 'Hot_Stars']
+
+        if stellar_type in easy_types:
+
+            PP = torch.as_tensor(self.polynomial_powers[stellar_type], dtype = torch.float64)
+            X  = torch.prod(K**PP, dim = -1)
+
+            flux = np.exp(self.coefficients[stellar_type] @ X)
+            flux *= self.reference[stellar_type]
+
+        elif stellar_type is 'Hottish_Giants':
+
+            PP = torch.as_tensor(self.polynomial_powers['Warm_Giants'], dtype = torch.float64)
+            X  = torch.prod(K**PP, dim = -1)
+
+            flux1 = np.exp(self.coefficients['Warm_Giants'] @ X)
+            flux1 *= self.reference['Warm_Giants']
+
+            PP = torch.as_tensor(self.polynomial_powers['Hot_Stars'], dtype = torch.float64)
+            X  = torch.prod(K**PP, dim = -1)
+
+            flux2 = np.exp(self.coefficients['Hot_Stars'] @ X)
+            flux2 *= self.reference['Hot_Stars']
+
+            t_index = (np.abs(gh_teff_overlap - teff2)).argmin()
+            weight = gh_weights[t_index]
+            flux = (flux1*weight + flux2*(1-weight))
+
+        elif stellar_type is 'Coolish_Giants':
+
+            PP = torch.as_tensor(self.polynomial_powers['Warm_Giants'], dtype = torch.float64)
+            X  = torch.prod(K**PP, dim = -1)
+
+            flux1 = np.exp(self.coefficients['Warm_Giants'] @ X)
+            flux1 *= self.reference['Warm_Giants']
+
+            PP = torch.as_tensor(self.polynomial_powers['Cool_Giants'], dtype = torch.float64)
+            X  = torch.prod(K**PP, dim = -1)
+
+            flux2 = np.exp(self.coefficients['Cool_Giants'] @ X)
+            flux2 *= self.reference['Cool_Giants']
+
+            t_index = (np.abs(gh_teff_overlap - teff2)).argmin()
+            weight = gc_weights[t_index]
+            flux = (flux1*weight + flux2*(1-weight))
+
+        elif stellar_type is 'Coolish_Dwarfs':
+
+            PP = torch.as_tensor(self.polynomial_powers['Warm_Dwarfs'], dtype = torch.float64)
+            X  = torch.prod(K**PP, dim = -1)
+
+            flux1 = np.exp(self.coefficients['Warm_Dwarfs'] @ X)
+            flux1 *= self.reference['Warm_Dwarfs']
+
+            PP = torch.as_tensor(self.polynomial_powers['Cool_Dwarfs'], dtype = torch.float64)
+            X  = torch.prod(K**PP, dim = -1)
+
+            flux2 = np.exp(self.coefficients['Cool_Dwarfs'] @ X)
+            flux2 *= self.reference['Warm_Dwarfs']
+
+            t_index = (np.abs(gh_teff_overlap - teff2)).argmin()
+            weight = d_weights[t_index]
+            flux = (flux1*weight + flux2*(1-weight))
+
+        else:
+            error = ('Parameter out of bounds:'
+                     'teff = {0},  logg {1}')
+            raise ValueError(error.format(teff2, logg))
+
 
         return flux