Basic ssp dev

lighthouse/SSP/basic_ssp.py

@@ -3,61 +3,148 @@
 from time import process_time as time
 
 import matplotlib.pyplot as plt
+import numpy as np
+from scipy.integrate import quad
 
 from .. import utils
 
 class Basic_SSP():
 
     def __init__(
             self,
-            isochrone: "Isochrone",
+            isochrone_model: "Isochrone",
             imf: "Initial_Mass_Function",
             sas: "Stellar_Atmosphere_Spectrum",
     ):
 
-        self.isochrone = isochrone
+        self.isochrone_grid = isochrone_model
         self.imf = imf
         self.sas = sas
 
-    def forward(self, metalicity, Tage, alpha) -> torch.Tensor:
 
-        isochrone = self.isochrone.get_isochrone(metalicity, Tage)
+        self.isochrone    = None
+        self.spectrum     = None
+        self._imf_weights = None
+        self._imf_weights_v2 = None
+
+
+    @property
+    def imf_weights_v2(self) -> torch.Tensor:
+
+        # like alf
+
+        if self._imf_weights_v2 is None:
+
+            weights = self.imf.get_imf(self.isochrone["initial_mass"], mass_weighted=False)
+            weights = weights/self.imf.t0_normalization
+
+            self._imf_weights_v2 = weights
+
+        return self._imf_weights_v2
+
+
+
+    @property
+    def imf_weights(self) -> torch.Tensor:
+
+        # like fsps
+
+
+        if self._imf_weights is None:
+
+            initial_stellar_masses = self.isochrone["initial_mass"]
+
+            lower_limit = self.imf.lower_limit
+            upper_limit = self.imf.upper_limit
+
+            weights = torch.zeros(initial_stellar_masses.shape)
+            for i, mass in enumerate(initial_stellar_masses):
+                if initial_stellar_masses[i] < lower_limit or initial_stellar_masses[i] > upper_limit:
+                    print("Bounds of isochrone exceed limits of full IMF")
+                if i == 0:
+                    m1 = lower_limit # ala fsps aka
+                else:
+                    m1 = initial_stellar_masses[i] - 0.5*(initial_stellar_masses[i] - initial_stellar_masses[i-1])
+                if i == len(initial_stellar_masses) - 1:
+                    m2 = initial_stellar_masses[i]
+                else:
+                    m2 = initial_stellar_masses[i] + 0.5*(initial_stellar_masses[i+1] - initial_stellar_masses[i])
+
+                if m2 < m1:
+                    print("IMF_WEIGHT WARNING: non-monotonic mass!", m1, m2, m2-m1)
+                    continue
+
+                if m2 == m1:
+                    print("m2 == m1")
+                    continue
+
+                weights[i], error =  quad(self.imf.get_imf,
+                                    m1, m2,
+                                    args=(False,) ) # i.e., not mass-weighting
+
+            self._imf_weights = weights/self.imf.t0_normalization
+
+        return self._imf_weights
+
+
+
+    def spectral_synthesis(self, metalicity, Tage) -> torch.Tensor:
+
+        isochrone   = self.isochrone
+        imf_weights = self.imf_weights_v2
+
+
+        ## https://waps.cfa.harvard.edu/MIST/README_tables.pdf
+        ## FSPS phase type defined as follows:
+        ## -1=PMS, 0=MS, 2=RGB, 3=CHeB, 4=EAGB,
+        ##  5=TPAGB, 6=postAGB, 9=WR
 
         # Main Sequence isochrone integration
-        CHOOSE = isochrone["phase"] <= 2
-        spectra = torch.stack(tuple(
-            self.sas.get_spectrum(
-               tf,
-               lg,
-                metalicity,
-            ) for lg, tf in zip(isochrone["log_g"][CHOOSE], isochrone["Teff"][CHOOSE])
+        MS = torch.logical_and(isochrone["phase"] >= 0, isochrone["phase"] <= 2)
+        # Horizontal Branch isochrone integration
+        HB = torch.logical_and(isochrone["phase"] > 2, isochrone["phase"] <= 5)
+
+
+        ms_spectra = torch.stack(tuple(
+            self.sas.get_spectrum(tf, lg, metalicity,
+            ) for lg, tf in zip(isochrone["log_g"][MS], isochrone["Teff"][MS])
         )).T
-        spectrum = torch.zeros(spectra.shape[0])
-        spectrum += torch.vmap(partial(torch.trapz, x = isochrone["initial_mass"][CHOOSE]))(
-            (spectra * self.imf.get_weight(
-                isochrone["initial_mass"][CHOOSE],
-                alpha,
-            ) * 10**isochrone["log_l"][CHOOSE]),
-        )
+        ms_spectra *= 10**isochrone["log_l"][MS]
 
-        # Horizontal Branch isochrone integration
-        CHOOSE = torch.logical_and(isochrone["phase"] > 2, isochrone["phase"] <= 5)
-        spectra = torch.stack(tuple(
-            self.sas.get_spectrum(
-                tf,
-                lg,
-                metalicity,
-            ) for lg, tf in zip(isochrone["log_g"][CHOOSE], isochrone["Teff"][CHOOSE])
+        hb_spectra = torch.stack(tuple(
+            self.sas.get_spectrum(tf, lg, metalicity,
+            ) for lg, tf in zip(isochrone["log_g"][HB], isochrone["Teff"][HB])
         )).T
-        spectrum += torch.vmap(partial(torch.trapz, x = isochrone["initial_mass"][CHOOSE]))(
-            (spectra * self.imf.get_weight(
-                isochrone["initial_mass"][CHOOSE],
-                alpha,
-            ) * 10**isochrone["log_l"][CHOOSE]),
-        )
+        hb_spectra *= 10**isochrone["log_l"][HB]
+
+        spectrum = torch.zeros(ms_spectra.shape[0])
+
+        spectrum += torch.vmap(partial(torch.trapz, x = isochrone["initial_mass"][MS])) (
+                                imf_weights[MS]*ms_spectra
+                    )
+        spectrum += torch.vmap(partial(torch.trapz, x = isochrone["initial_mass"][HB])) (
+                                imf_weights[HB]*hb_spectra
+                    )
+
+
 
         # SSP in L_sun Hz^-1, CvD models in L_sun micron^-1, convert
-        spectrum *= utils.light_speed/self.sas.wavelength**2
+        spectrum *= utils.light_speed_micron/self.sas.wavelength**2 #for detailed comparisons to alf
+        #spectrum *= utils.light_speed_cgs/self.sas.wavelength**2 # the correct thing to do;consistent normalization with fsps
+
+        return spectrum
+
+
+    def forward(self, metalicity, Tage, synthesize_spectrum=True) -> torch.Tensor:
+
+        self.isochrone = self.isochrone_grid.get_isochrone(metalicity, Tage)
+
+        if synthesize_spectrum:
+            spectrum = self.spectral_synthesis(metalicity, Tage)
+        else:
+            print("HAVE NOT SYNTHEISIZED A SPECTRUM")
+            spectrum = None
+
 
         return spectrum
 

lighthouse/initial_mass_function/__init__.py

@@ -1,3 +1,5 @@
 from .kroupa import *
+from .salpeter import *
+from .two_slope_powerlaw import *
 from .initial_mass_function import *
 

lighthouse/initial_mass_function/initial_mass_function.py

@@ -1,14 +1,38 @@
 from abc import ABC, abstractmethod
-from torch import Tensor
+# from torch import Tensor
+
+import torch
+
+from scipy.integrate import quad
+import numpy as np
+
 
 __all__ = ("Initial_Mass_Function", )
 
 class Initial_Mass_Function(ABC):
 
-    @abstractmethod
-    def get_weight(self, mass) -> Tensor:
-        ...
+    def __init__(self):
+
+        self.lower_limit = torch.tensor(0.08, dtype=torch.float64)
+        self.upper_limit = torch.tensor(100., dtype=torch.float64)
+
+
+        self._t0_normalization = None
+
 
     @abstractmethod
-    def to(self, dtype=None, device=None):
-        ...
+    def get_imf(self, mass, mass_weighted=False) -> torch.Tensor:
+        pass
+
+
+    @property
+    def t0_normalization(self): ## TODO: change this name
+
+        if self._t0_normalization is None:
+            self._t0_normalization = quad(self.get_imf,
+                        self.lower_limit,
+                        self.upper_limit,
+                        args=(True,) )[0]
+
+        return self._t0_normalization
+

lighthouse/initial_mass_function/kroupa.py

@@ -1,14 +1,18 @@
 import torch
 
+
 from .initial_mass_function import Initial_Mass_Function
 
 __all__ = ("Kroupa", )
 
 class Kroupa(Initial_Mass_Function):
 
-    def get_weight(self, mass, alpha) -> torch.Tensor:
+    def get_imf(self, mass, mass_weighted=False) -> torch.Tensor:
+
+        alpha = torch.tensor([1.3, 2.3, 2.3])
+        mass  = torch.tensor(mass)
 
-        weight = torch.where(
+        imf = torch.where(
             mass < 0.5,
             mass**(-alpha[0]), # mass < 0.5
             torch.where(
@@ -18,8 +22,12 @@ def get_weight(self, mass, alpha) -> torch.Tensor:
             )
         )
 
-        return weight
+        if mass_weighted:
+            return mass*imf
+        else:
+            return imf
 
+
     def to(self, dtype=None, device=None):
         pass
 
@@ -30,8 +38,8 @@ def to(self, dtype=None, device=None):
     K = Kroupa()
 
     M = torch.linspace(0.1, 100, 1000)
-    W = K.get_weight(M, torch.tensor([1.3, 2.3, 2.7]))
+    IMF = K.get_imf(M)
 
-    plt.plot(torch.log10(M), torch.log10(W))
+    plt.plot(torch.log10(M), torch.log10(IMF))
     plt.show()
 

lighthouse/initial_mass_function/salpeter.py

@@ -0,0 +1,33 @@
+import torch
+
+
+
+
+from .initial_mass_function import Initial_Mass_Function
+
+__all__ = ("Salpeter", )
+
+class Salpeter(Initial_Mass_Function):
+
+    def get_imf(self, mass, mass_weighted=False) -> torch.Tensor:
+
+        salpeter_index = torch.tensor(2.35, dtype = torch.float64)
+        imf = mass**(-salpeter_index)
+
+        if mass_weighted:
+            imf = imf*mass
+
+        return imf
+
+
+if __name__ == "__main__":
+
+    import matplotlib.pyplot as plt
+    K = Salpeter()
+
+    M = torch.linspace(0.08, 100, 1000)
+    IMF = K.get_imf(M)
+
+    plt.plot(torch.log10(M), torch.log10(IMF))
+    plt.show()
+

lighthouse/initial_mass_function/two_slope_powerlaw.py

@@ -0,0 +1,38 @@
+import torch
+
+
+from .initial_mass_function import Initial_Mass_Function
+
+__all__ = ("Two_Slope_Powerlaw", )
+
+class Two_Slope_Powerlaw(Initial_Mass_Function):
+
+    """
+    Fixed lower mass cutoff
+    Fixed break 
+    """
+
+    def get_imf(self, mass, mass_weighted=False, alpha1=1.0, alpah2=1.0) -> torch.Tensor:
+
+        alpha = torch.tensor([alpha1, alpha2, -2.30])
+        mass  = torch.tensor(mass)
+
+        imf = torch.where(
+            mass < 0.5,
+            mass**(-alpha[0]), # mass < 0.5
+            torch.where(
+                mass < 1.0,
+                0.5**(-alpha[0] + alpha[1]) * mass**(-alpha[1]), # 0.5 <= mass < 1.0
+                0.5**(-alpha[0] + alpha[1]) * mass**(-alpha[2]), # mass >= 1.0
+            )
+        )
+
+        if mass_weighted:
+            return mass*imf
+        else:
+            return imf
+
+
+if __name__ == "__main__":
+
+    pass