Merge pull request #13 from kadrlica/imf

Revamp IMF module (closes #15).

tests/test_imf.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python
+"""
+Generic python script.
+"""
+__author__ = "Alex Drlica-Wagner"
+import numpy as np
+
+import ugali.analysis.imf
+import ugali.isochrone
+
+MASSES = np.array([0.1, 0.5, 1.0, 10.0])
+SEED = 1
+IMFS = ['Salpeter1955','Kroupa2001','Chabrier2003']
+
+def test_imf():
+    dn_dm = ugali.analysis.imf.chabrierIMF(MASSES)
+    np.testing.assert_allclose(dn_dm, 
+                               [1.29919665, 0.66922584, 0.3666024, 0.01837364],
+                               rtol=1e-6)
+
+    #imf = ugali.analysis.imf.IMF('chabrier')
+    imf = ugali.analysis.imf.factory('Chabrier2003')
+    masses = imf.sample(3, seed=SEED)
+    np.testing.assert_allclose(masses,[0.22122041, 0.4878909, 0.10002029],rtol=1e-6)
+
+    imf = ugali.analysis.imf.Chabrier2003()
+    masses = imf.sample(3, seed=SEED)
+    np.testing.assert_allclose(masses,[0.22122041, 0.4878909, 0.10002029],rtol=1e-6)
+
+    iso = ugali.isochrone.Bressan2012(imf_type='Chabrier2003')
+    masses = iso.imf.sample(3,seed=SEED)
+    np.testing.assert_allclose(masses,[0.22122041, 0.4878909, 0.10002029],rtol=1e-6)
+
+    integral = iso.imf.integrate(0.1,2.0)
+    np.testing.assert_allclose(integral,0.94961632593)
+
+def test_imf_norm():
+    # All of the IMFs have normalized integrals from 0.1 to 100 Msun
+    for n in IMFS:
+        print(n)
+        imf = ugali.analysis.imf.factory(name=n)
+        np.testing.assert_allclose(imf.integrate(0.1,100,steps=int(1e4)),
+                                   1.0,rtol=1e-3)
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description=__doc__)
+    args = parser.parse_args()
+
+    #test_imf()
+    test_imf_norm()

tests/test_isochrone.py

@@ -142,9 +142,9 @@ def test_download():
         try:
             iso.download(outdir='./tmp/'+name.lower(),force=True)
         except URLError as e:
-            print("Server is down")
-            print(e)
-
+            logger.error("Server is down:\n%s"%str(e))
+        except RuntimeError as e:
+            logger.error(str(e))
 
 if __name__ == "__main__":
     import argparse

ugali/analysis/imf.py

@@ -1,47 +1,39 @@
 """
 Classes to handle initial mass functions (IMFs).
-"""
 
-import numpy
+https://github.com/keflavich/imf
+"""
+from abc import abstractmethod
 import numpy as np
 import scipy.interpolate
 
 from ugali.utils.logger import logger
 
-# ADW: TODO - This needs to be modernized
-#      Also add Kroupa and Salpeter IMFs...
-
 ############################################################
 
-class IMF:
+class IMF(object):
     """
     Base class for initial mass functions (IMFs).
     """
 
-    def __init__(self, type='chabrier'):
-        """
-        Initialize an instance of an initial mass function.
-        """
+    def __call__(self, mass, **kwargs):
+        """ Call the pdf of the mass function """
+        return self.pdf(mass,**kwargs)
 
-        self.type = type
+    def integrate(self, mass_min, mass_max, log_mode=True, weight=False, steps=1e4):
+        """ Numerical Riemannn integral of the IMF (stupid simple).
 
-        if self.type == 'chabrier':
-            self.pdf = chabrierIMF
-        else:
-            logger.warn('initial mass function type %s not recognized'%(self.type))
-
-    def integrate(self, mass_min, mass_max, log_mode=True, weight=False, steps=10000):
-        """
-        Numerically integrate initial mass function.
-
-        INPUTS:
-            mass_min: minimum mass bound for integration (solar masses)
-            mass_max: maximum mass bound for integration (solar masses)
-            log_mode[True]: use logarithmic steps in stellar mass as oppose to linear
+        Parameters:
+        -----------
+        mass_min: minimum mass bound for integration (solar masses)
+        mass_max: maximum mass bound for integration (solar masses)
+        log_mode[True]: use logarithmic steps in stellar mass as oppose to linear
             weight[False]: weight the integral by stellar mass
             steps: number of numerical integration steps
-        OUTPUT:
-            result of integral
+
+        Returns:
+        --------
+        result of integral
         """
         if log_mode:
             d_log_mass = (np.log10(mass_max) - np.log10(mass_min)) / float(steps)
@@ -63,59 +55,166 @@ def integrate(self, mass_min, mass_max, log_mode=True, weight=False, steps=10000
 
     def sample(self, n, mass_min=0.1, mass_max=10., steps=10000, seed=None):
         """
-        Sample n initial mass values between mass_min and mass_max, following the IMF distribution.
+        Sample initial mass values between mass_min and mass_max,
+        following the IMF distribution.
+
+        ADW: Should this be `sample` or `simulate`?
+
+        Parameters:
+        -----------
+        n : number of samples to draw
+        mass_min : minimum mass to sample from
+        mass_max : maximum mass to sample from
+        steps    : number of steps for isochrone sampling
+        seed     : random seed (passed to np.random.seed)
+
+        Returns:
+        --------
+        mass     : array of randomly sampled mass values
         """
         if seed is not None: np.random.seed(seed)
         d_mass = (mass_max - mass_min) / float(steps)
-        mass = numpy.linspace(mass_min, mass_max, steps)
-        cdf = numpy.insert(numpy.cumsum(d_mass * self.pdf(mass[1:], log_mode=False)), 0, 0.)
+        mass = np.linspace(mass_min, mass_max, steps)
+        cdf = np.insert(np.cumsum(d_mass * self.pdf(mass[1:], log_mode=False)), 0, 0.)
         cdf = cdf / cdf[-1]
         f = scipy.interpolate.interp1d(cdf, mass)
-        return f(numpy.random.uniform(size=n))
-
-############################################################
+        return f(np.random.uniform(size=n))
 
-def chabrierIMF(mass, log_mode=True, a=1.31357499301):
-    """
-    Chabrier initial mass function. 
+    @abstractmethod
+    def pdf(cls, mass, **kwargs): pass
+        
 
+class Chabrier2003(IMF):
+    """ Initial mass function from Chabrier (2003):
     "Galactic Stellar and Substellar Initial Mass Function"
     Chabrier PASP 115:763-796 (2003)
     https://arxiv.org/abs/astro-ph/0304382    
+    """
+
+    @classmethod
+    def pdf(cls, mass, log_mode=True):
+        """ PDF for the Chabrier IMF.
+         
+        The functional form and coefficients are described in Eq 17
+        and Tab 1 of Chabrier (2003):
+
+          m <= 1 Msun: E(log m) = A1*exp(-(log m - log m_c)^2 / 2 sigma^2)
+          m  > 1 Msun: E(log m) = A2 * m^-x
+         
+          A1 = 1.58 : normalization [ log(Msun)^-1 pc^-3]
+          m_c = 0.079 [Msun]
+          sigma = 0.69
+          A2 = 4.43e-2
+          x = 1.3
+          
+        We redefine a = A1, A2 = a * b;
+
+        The normalization is set so that the IMF integrates to 1 over
+        the mass range from 0.1 Msun to 100 Msun
+
+        Parameters:
+        -----------
+        mass: stellar mass (solar masses)
+        log_mode[True]: return number per logarithmic mass range, i.e., dN/dlog(M)
+         
+        Returns:
+        --------
+        number per (linear or log) mass bin, i.e., dN/dM or dN/dlog(M) where mass unit is solar masses
+        """
+        log_mass = np.log10(mass)
+
+        # Constants from Chabrier 2003
+        m_c = 0.079
+        sigma = 0.69
+        x = 1.3 
+
+        # This value is set to normalize integral from 0.1 to 100 Msun
+        a=1.31357499301
+        # This value is required so that the two components match at 1 Msun
+        b = 0.279087531047 
+
+        dn_dlogm = ((log_mass <= 0) * a * np.exp(-(log_mass - np.log10(m_c))**2 / (2 * (sigma**2)))) 
+        dn_dlogm += ((log_mass  > 0) * a * b * mass**(-x))
+
+        if log_mode:
+            # Number per logarithmic mass range, i.e., dN/dlog(M)
+            return dn_dlogm
+        else:
+            # Number per linear mass range, i.e., dN/dM
+            return dn_dlogm / (mass * np.log(10))
+
+class Kroupa2001(IMF):
+    """ IMF from Kroupa (2001):
 
-    The form and coefficients are described in Equation 17 and Table 1:
-      m <= 1 Msun: E(log m) = A1*exp(-(log m - log m_c)^2 / 2 sigma^2)
-      m  > 1 Msun: E(log m) = A2 * m^-x
-
-      A1 = 1.58 : normalization [ log(Msun)^-1 pc^-3]
-      m_c = 0.079 [Msun]
-      sigma = 0.69
-      A2 = 4.43e-2
-      x = 1.3
-      
-    We redefine a = A1, A2 = a * b;
-      
-    Chabrier set's his normalization 
-
-    Parameters:
-    -----------
-    mass: stellar mass (solar masses)
-    log_mode[True]: return number per logarithmic mass range, i.e., dN/dlog(M)
-    a[1.31357499301]: normalization; normalized by default to the mass interval 0.1--100 solar masses
+    "On the variation of the initial mass function"
+    MNRAS 322:231 (2001)
+    https://arxiv.org/abs/astro-ph/0009005
+    """
+
+    @classmethod
+    def pdf(cls, mass, log_mode=True):
+        """ PDF for the Kroupa IMF.
+
+        Normalization is set over the mass range from 0.1 Msun to 100 Msun
+        """
+        log_mass = np.log10(mass)
+        # From Eq 2
+        mb = mbreak  = [0.08, 0.5] # Msun
+        a = alpha = [0.3, 1.3, 2.3] # alpha
+
+        # Normalization set from 0.1 -- 100 Msun
+        norm = 0.27947743949440446
+
+        b = 1./norm
+        c = b * mbreak[0]**(alpha[1]-alpha[0])
+        d = c * mbreak[1]**(alpha[2]-alpha[1])
+
+        dn_dm  = b * (mass < 0.08) * mass**(-alpha[0])
+        dn_dm += c * (0.08 <= mass) * (mass < 0.5) * mass**(-alpha[1])
+        dn_dm += d * (0.5  <= mass) * mass**(-alpha[2])
+
+        if log_mode:
+            # Number per logarithmic mass range, i.e., dN/dlog(M)
+            return dn_dm * (mass * np.log(10))
+        else:
+            # Number per linear mass range, i.e., dN/dM
+            return dn_dm
 
-    Returns:
-    --------
-    number per (linear or log) mass bin, i.e., dN/dM or dN/dlog(M) where mass unit is solar masses
+class Salpeter1955(IMF): 
+    """ IMF from Salpeter (1955):
+    "The Luminosity Function and Stellar Evolution"
+    ApJ 121, 161S (1955)
+    http://adsabs.harvard.edu/abs/1955ApJ...121..161S
     """
-    log_mass = np.log10(mass)
-    b = 0.279087531047 # This value is required so that the two components match at 1 Msun
-    if log_mode:
-        # Number per logarithmic mass range, i.e., dN/dlog(M)
-        return ((log_mass <= 0) * a * np.exp(-(log_mass - np.log10(0.079))**2 / (2 * (0.69**2)))) + \
-               ((log_mass  > 0) * a * b * mass**(-1.3))
-    else:
-        # Number per linear mass range, i.e., dN/dM
-        return (((log_mass <= 0) * a * np.exp(-(log_mass - np.log10(0.079))**2 / (2 * (0.69**2)))) + \
-                ((log_mass  > 0) * a * b * mass**(-1.3))) / (mass * np.log(10))
+
+    @classmethod
+    def pdf(cls, mass, log_mode=True):
+        """ PDF for the Salpeter IMF.
+
+        Value of 'a' is set to normalize the IMF to 1 between 0.1 and 100 Msun
+        """
+        alpha = 2.35
+
+        a = 0.060285569480482866
+        dn_dm  = a * mass**(-alpha)
+
+        if log_mode:
+            # Number per logarithmic mass range, i.e., dN/dlog(M)
+            return dn_dm * (mass * np.log(10))
+        else:
+            # Number per linear mass range, i.e., dN/dM
+            return dn_dm
+
+def factory(name, **kwargs):
+    from ugali.utils.factory import factory
+    return factory(name, module=__name__, **kwargs)
+
+imfFactory = factory
+
+############################################################
+
+def chabrierIMF(mass, log_mode=True):
+    """ Backward compatible wrapper around Chabrier2003.pdf """
+    return Chabrier2003.pdf(mass,log_mode=log_mode)
 
 ############################################################

ugali/config/config_dr10_gal.yaml

@@ -179,7 +179,7 @@ isochrone:
   mag_1_field: g
   mag_2_field: r
   stage_field: stage
-  imf: chabrier
+  imf: Chabrier2003
   horizontal_branch_stage     : 'BHeb'
   #horizontal_branch_dispersion: null
   #horizontal_branch_dispersion: [0.0] # mag

ugali/config/config_sva1.yaml

@@ -172,7 +172,7 @@ isochrone:
   mag_1_field: g
   mag_2_field: r
   stage_field: stage
-  imf: chabrier
+  imf: Chabrier2003
   horizontal_branch_stage     : 'BHeb'
   #horizontal_branch_dispersion: null
   #horizontal_branch_dispersion: [0.0] # mag

ugali/isochrone/model.py

@@ -84,7 +84,7 @@ class IsochroneModel(Model):
         ('band_1','g','Field name for magnitude one'),
         ('band_2','r','Field name for magnitude two'),
         ('band_1_detection',True,'Band one is detection band'),
-        ('imf_type','chabrier','Initial mass function'),
+        ('imf_type','Chabrier2003','Initial mass function'),
         ('hb_stage',None,'Horizontal branch stage name'),
         ('hb_spread',0.0,'Intrinisic spread added to horizontal branch'),
         )
@@ -100,7 +100,7 @@ def _setup(self, **kwargs):
         for k,v in list(defaults.items()):
             setattr(self,k,v)
 
-        self.imf = ugali.analysis.imf.IMF(defaults['imf_type'])
+        self.imf = ugali.analysis.imf.factory(defaults['imf_type'])
         self.index = None
 
     def _parse(self,filename):