Merge pull request #49 from xuanxu/new-imf

Add Kroupa 2001 IMF

Author: xuanxu

README.rst

@@ -68,7 +68,7 @@ Starmatrix reads a config file where several options can be set in yaml format::
 
         z: 0.0200               # metallicity
         sol_ab: as09            # solar abundances
-        imf: kroupa             # initial mass function (IMF)
+        imf: kroupa2002         # initial mass function (IMF)
         imf_m_low: 0.15         # lower mass limit for the IMF
         imf_m_up: 100           # upper mass limit for the IMF
         total_time_steps: 300   # number of time steps (will result in a Q Matrix per step)
@@ -99,11 +99,12 @@ The ``imf`` param in the config file can be set to use any of the predefined IMF
 :starburst: Starburst 1999 (a Salpeter with mass limits in [1, 120])
 :miller_scalo: Miller & Scalo 1979
 :ferrini: Ferrini, Palla & Penco 1998
-:kroupa: Kroupa 2002
+:kroupa2001: Kroupa 2001
+:kroupa2002: Kroupa 2002
 :chabrier: Chabrier 2003
 :maschberger: Maschberger 2012
 
-The default value is ``kroupa``. If you want to use your own IMF you can do so subclassing the `IMF class`_.
+The default value is ``kroupa2002``. If you want to use your own IMF you can do so subclassing the `IMF class`_.
 
 .. _`IMF class`: https://github.com/xuanxu/starmatrix/blob/main/src/starmatrix/imfs.py#L35-L68
 

docs/configuration.rst

@@ -44,11 +44,12 @@ The ``imf`` param in the config file can be set to use any of the predefined IMF
 :starburst: Starburst 1999 (a Salpeter with mass limits in [1, 120])
 :miller_scalo: Miller & Scalo 1979
 :ferrini: Ferrini, Palla & Penco 1998
-:kroupa: Kroupa 2002
+:kroupa2001: Kroupa 2001
+:kroupa2002: Kroupa 2002
 :chabrier: Chabrier 2003
 :maschberger: Maschberger 2012
 
-The default value is ``kroupa``. If you want to use your own IMF you can do so subclassing the `IMF class`_.
+The default value is ``kroupa2002``. If you want to use your own IMF you can do so subclassing the `IMF class`_.
 
 .. _`IMF class`: https://github.com/xuanxu/starmatrix/blob/main/src/starmatrix/imfs.py#L20-L40
 

src/starmatrix/imfs.py

@@ -7,7 +7,7 @@
 * Miller & Scalo 1979
 * Ferrini, Palla & Penco 1998
 * Starburst 1999
-* Kroupa 2002
+* Kroupa 2001 & 2002
 * Chabrier 2003
 * Maschberger 2012
 
@@ -25,7 +25,8 @@ def select_imf(name, params={}):
         "starburst": Starburst,
         "chabrier": Chabrier,
         "ferrini": Ferrini,
-        "kroupa": Kroupa,
+        "kroupa2002": Kroupa2002,
+        "kroupa2001": Kroupa2001,
         "miller_scalo": MillerScalo,
         "maschberger": Maschberger
     }
@@ -120,7 +121,7 @@ def description(self):
         return "IMF Ferrini, Palla & Penco 1998"
 
 
-class Kroupa(IMF):
+class Kroupa2002(IMF):
     def m_phi(self, m):
         if 0.015 <= m < 0.08:
             return m * (m ** -0.35)
@@ -137,6 +138,21 @@ def description(self):
         return "IMF from Kroupa 2002"
 
 
+class Kroupa2001(IMF):
+    def m_phi(self, m):
+        if 0.015 <= m < 0.08:
+            return m * (m ** -0.35)
+        elif 0.08 <= m < 0.5:
+            return m * 0.08 * (m ** -1.3)
+        elif 0.5 <= m:
+            return m * 0.04 * (m ** -2.3)
+        else:
+            return 0
+
+    def description(self):
+        return "IMF from Kroupa 2001"
+
+
 class Chabrier(IMF):
     def m_phi(self, m):
         if m <= 1:

src/starmatrix/sample_input/params.yml

@@ -1,7 +1,7 @@
 # All configurable parameters:
 # z                           -> Metallicity. Default value: 0.02
 # sol_ab                      -> Solar abundances data (*). Default value: as09
-# imf                         -> Initial Mass function to use (*). Default value: kroupa
+# imf                         -> Initial Mass function to use (*). Default value: kroupa2002
 # imf_alpha                   -> If IMF is salpeter/starburst, this extra param is needed. Defaults to 2.35
 # imf_m_low                   -> Lower limit (in solar masses) for the IMF. Default value: 0.15
 # imf_m_up                    -> Upper limit (in solar masses) for the IMF. Default value: 100
@@ -28,7 +28,8 @@
 #         starburst    = Starburst 1999 (special case of Salpeter in [1, 120])
 #         chabrier     = Chabrier 2003
 #         ferrini      = Ferrini, Palla & Penco 1998
-#         kroupa       = Kroupa 2002 (default)
+#         kroupa2001   = Kroupa 2001
+#         kroupa2002   = Kroupa 2002 (default)
 #         miller_scalo = Miller & Scalo 1979
 #         maschberger  = Maschberger 2012
 #       ]
@@ -79,7 +80,7 @@
 
 z: 0.02
 sol_ab: as09
-imf: kroupa
+imf: kroupa2002
 m_max: 40.0
 dtd_sn: rlp
 expelled_elements_filename: ejections.dat

src/starmatrix/settings.py

@@ -12,7 +12,7 @@
 default = {
     "z": 0.02,
     "sol_ab": "as09",
-    "imf": "kroupa",
+    "imf": "kroupa2002",
     "imf_alpha": 2.35,
     "imf_m_low": 0.15,
     "imf_m_up": 100,
@@ -32,7 +32,7 @@
 }
 
 valid_values = {
-    "imf": ["salpeter", "starburst", "chabrier", "ferrini", "kroupa", "miller_scalo", "maschberger"],
+    "imf": ["salpeter", "starburst", "chabrier", "ferrini", "kroupa2001", "kroupa2002", "miller_scalo", "maschberger"],
     "dtd_sn": ["rlp", "maoz", "castrillo", "greggio", "chen", "greggio-CDD04", "greggio-CDD1",
                "greggio-WDD04", "greggio-WDD1", "greggio-SDCH", "greggio-SDSCH",
                "strolger-fit1", "strolger-fit2", "strolger-fit3", "strolger-fit4", "strolger-fit5", "strolger-optimized"],

src/starmatrix/tests/test_imfs.py

@@ -3,7 +3,7 @@
 import numpy as np
 import starmatrix.settings as settings
 from starmatrix.imfs import select_imf, IMF
-from starmatrix.imfs import Salpeter, Starburst, Chabrier, Ferrini, Kroupa, MillerScalo, Maschberger
+from starmatrix.imfs import Salpeter, Starburst, Chabrier, Ferrini, Kroupa2001, Kroupa2002, MillerScalo, Maschberger
 
 
 @pytest.fixture
@@ -15,8 +15,8 @@ def available_imfs():
 
 
 def test_select_imf():
-    strings = ["salpeter", "starburst", "chabrier", "ferrini", "kroupa", "miller_scalo", "maschberger"]
-    classes = [Salpeter, Starburst, Chabrier, Ferrini, Kroupa, MillerScalo, Maschberger]
+    strings = ["salpeter", "starburst", "chabrier", "ferrini", "kroupa2001", "kroupa2002", "miller_scalo", "maschberger"]
+    classes = [Salpeter, Starburst, Chabrier, Ferrini, Kroupa2001, Kroupa2002, MillerScalo, Maschberger]
 
     for i in range(len(strings)):
         imf_instance = select_imf(strings[i])
@@ -47,9 +47,11 @@ def test_imf_is_zero_if_no_positive_mass(available_imfs):
             assert select_imf(imf).for_mass(mass) > 0.0
 
 
-def test_minimum_mass_value_for_kroupa_imf():
-    assert select_imf("kroupa").for_mass(0.014) == 0.0
-    assert select_imf("kroupa").for_mass(0.015) > 0.0
+def test_minimum_mass_value_for_kroupa_imfs():
+    assert select_imf("kroupa2001").for_mass(0.014) == 0.0
+    assert select_imf("kroupa2001").for_mass(0.015) > 0.0
+    assert select_imf("kroupa2002").for_mass(0.014) == 0.0
+    assert select_imf("kroupa2002").for_mass(0.015) > 0.0
 
 
 def test_for_mass_is_normalized(available_imfs):