Add the initial mass function scripts

.gitlab-ci.yml

@@ -35,6 +35,5 @@ before_script:
 
 test:
   script:
-    - cd test
-    - ./getCoolingTable.sh
-    - ./test_cooling.py
\ No newline at end of file
+    - cd tests
+    - ./run_tests.sh
\ No newline at end of file

examples/initial_mass_function/plot_initial_mass_function.py

+#!/usr/bin/env python3
+# ########################################################################
+# This file is part of PYSWIFT.
+# Copyright (c) 2018 Loic Hausammann (loic.hausammann@epfl.ch)
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+# ########################################################################
+
+from pyswiftsim import libstellar
+import pyswiftsim
+
+import numpy as np
+import matplotlib.pyplot as plt
+plt.rc('text', usetex=True)
+
+N = 100
+
+
+if __name__ == "__main__":
+
+    with pyswiftsim.ParameterFile() as filename:
+
+        parser = pyswiftsim.parseYamlFile(filename)
+        imf_parser = parser["GEARInitialMassFunction"]
+
+        mass_limits = list(imf_parser["mass_limits"])
+        mass_min = np.log10(mass_limits[0])
+        mass_max = np.log10(mass_limits[-1])
+
+        # du / dt
+        print("Computing initial mass function...")
+
+        mass = np.logspace(mass_min, mass_max, N, dtype=np.float32)
+        imf = libstellar.initialMassFunction(filename, mass)
+
+        # plot IMF
+        plt.loglog(mass, imf)
+
+        plt.xlabel("Mass [Msun]")
+        plt.ylabel("Mass fraction")
+
+        plt.show()

examples/initial_mass_function/run.sh

+#!/bin/bash
+
+./plot_initial_mass_function.py

pyswiftsim/__init__.py

@@ -32,6 +32,10 @@ def downloadCoolingTable():
         urllib.request.urlretrieve(url, filename)
 
 
+def downloadChemistryTable():
+    print("Not available")
+
+
 def parseYamlFile(filename):
     with open(filename, "r") as stream:
         stream = yaml.load(stream)
@@ -81,6 +85,18 @@ class ParameterFile:
 
         "GearChemistry": {
             "InitialMetallicity": 0.01295,
+        },
+
+        "GEARInitialMassFunction": {
+            "number_function_part": 4,
+            "exponents": [0.7, -0.8, -1.7, -1.3],
+            "mass_limits": [0.05, 0.08, 0.5, 1.0, 50.]
+        },
+
+        "GEARFeedback": {
+            "SupernovaeEnergy_erg": 1e51,
+            "ThermalTime_Myr": 5,
+            "ChemistryTable": "chemistry.hdf5",
         }
     }
 

scripts/pyswift_initial_mass_function

+#!/usr/bin/env python3
+# ########################################################################
+# This file is part of PYSWIFT.
+# Copyright (c) 2018 Loic Hausammann (loic.hausammann@epfl.ch)
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+# ########################################################################
+
+from pyswiftsim import libstellar
+import pyswiftsim
+
+import numpy as np
+import matplotlib.pyplot as plt
+from argparse import ArgumentParser
+import os
+
+plt.rc('text', usetex=True)
+
+
+def parseArguments():
+    parser = ArgumentParser(description="Plot the initial mass function")
+
+    parser.add_argument(
+        "--download_table", dest="table",
+        action="store_true",
+        help="Download the default chemistry table")
+
+    parser.add_argument(
+        "--params", dest="params",
+        type=str, default=None,
+        help="SWIFT parameters file")
+
+    parser.add_argument(
+        "--N", dest="N",
+        type=int, default=1000,
+        help="Number of points")
+
+    parser.add_argument(
+        "--mass_min", dest="mass_min",
+        type=float, default=None,
+        help="Minimal mass to plot")
+
+    parser.add_argument(
+        "--mass_max", dest="mass_max",
+        type=float, default=None,
+        help="Maximal mass to plot")
+
+    args = parser.parse_args()
+
+    return args
+
+
+if __name__ == "__main__":
+    opt = parseArguments()
+
+    if opt.table:
+        pyswiftsim.downloadChemistryTable()
+
+    if opt.params is not None and not os.path.isfile(opt.params):
+        raise Exception(
+            "The parameter file '{}' does not exist".format(opt.params))
+
+    with pyswiftsim.ParameterFile(existing=opt.params) as filename:
+
+        parser = pyswiftsim.parseYamlFile(filename)
+        imf_parser = parser["GEARInitialMassFunction"]
+
+        mass_limits = list(imf_parser["mass_limits"])
+        if opt.mass_min is not None:
+            mass_min = np.log10(opt.mass_min)
+        else:
+            mass_min = np.log10(mass_limits[0])
+
+        if opt.mass_max is not None:
+            mass_max = np.log10(opt.mass_max)
+        else:
+            mass_max = np.log10(mass_limits[-1])
+
+        # du / dt
+        print("Computing cooling...")
+
+        mass = np.logspace(mass_min, mass_max, opt.N, dtype=np.float32)
+        imf = libstellar.initialMassFunction(filename, mass)
+
+        # plot IMF
+        plt.loglog(mass, imf)
+
+        plt.xlabel("Mass [Msun]")
+        plt.ylabel("Mass fraction")
+
+        plt.show()

setup.py

@@ -21,10 +21,6 @@ cflags = [
     # disables some warnings due to python
     "-Wno-unused-parameter",
     "-Wno-maybe-uninitialized",
-    "-Wno-strict-prototypes",
-    "-Wno-unused-function",
-    "-Wno-incompatible-pointer-types",
-    "-Wno-missing-field-initializers",
     "-fopenmp",
 ]
 

src/cooling_wrapper.c

@@ -66,7 +66,7 @@ void pycooling_set_fractions(struct xpart *xp, PyArrayObject* frac, const int id
  * @brief Compute the cooling rate
  *
  * args is expecting:
- *   - a parameter filename,
+ *   - the name of the parameter file
  *   - a numpy array of the densities,
  *   - a numpy array of the specific energies,
  *   - an optional int if the cosmology should be used,

src/libstellar.c

@@ -27,7 +27,25 @@
 static PyMethodDef libstellar_methods[] = {
 
   {"initialMassFunction", pyinitial_mass_function, METH_VARARGS,
-   "TODO\n"
+   "Compute the initial mass function.\n\n"
+   "Parameters\n"
+   "----------\n\n"
+   "filename: str\n"
+   "\t Parameter file\n"
+   "mass: array\n"
+   "\t Mass of the stars in internal units\n"
+   "\n"
+   "Returns\n"
+   "-------\n"
+   "imf: array\n"
+   "\t The initial mass function\n"
+   "\n"
+   "Examples\n"
+   "--------\n"
+   ">>> mass = np.array([1, 10])\n"
+   ">>> filename = 'params.yml'\n"
+   ">>> imf = initialMassFunction(filename, mass)\n"
+
   },
 
   {"lifetimeFromMass", pylifetime_from_mass, METH_VARARGS,

src/stellar_evolution_wrapper.c

@@ -23,11 +23,12 @@
  * @brief Compute the initial mass function.
  *
  * args is expecting:
- * TODO
+ *  - the name of the parameter file
+ *  - a numpy array containing the masses.
  *
  * @param self calling object
  * @param args arguments
- * @return cooling rate
+ * @return The initial mass function
  */
 PyObject* pyinitial_mass_function(PyObject* self, PyObject* args) {
   import_array();
@@ -50,10 +51,10 @@ PyObject* pyinitial_mass_function(PyObject* self, PyObject* args) {
 
   int with_cosmo = 0;
   PYSWIFTSIM_INIT_STRUCTS();
-  
+
   /* Init stellar physics */
   struct feedback_props fp;
-  feedback_props_init(&fp, &pconst, &hydro_props, &us, &params, &cosmo);
+  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);
 
   /* return object */
   PyArrayObject *imf = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(mass), PyArray_DIMS(mass), NPY_FLOAT);
@@ -62,7 +63,7 @@ PyObject* pyinitial_mass_function(PyObject* self, PyObject* args) {
     float m = *(float*) PyArray_GETPTR1(mass, i);
     float *imf_i = (float*) PyArray_GETPTR1(imf, i);
 
-    
+
     *imf_i = stellar_evolution_get_imf(&fp.stellar_model.imf, m);
   }
 
@@ -111,7 +112,7 @@ PyObject* pylifetime_from_mass(PyObject* self, PyObject* args) {
   
   /* Init stellar physics */
   struct feedback_props fp;
-  feedback_props_init(&fp, &pconst, &hydro_props, &us, &params, &cosmo);
+  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);
 
   /* return object */
   PyArrayObject *lifetime = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(mass), PyArray_DIMS(mass), NPY_FLOAT);
@@ -122,7 +123,7 @@ PyObject* pylifetime_from_mass(PyObject* self, PyObject* args) {
     float *tau = (float*) PyArray_GETPTR1(lifetime, i);
 
     
-    *tau = stellar_evolution_get_log_lifetime_from_mass(&fp.stellar_model.imf, log10(m), z);
+    *tau = stellar_evolution_get_log_lifetime_from_mass(&fp.stellar_model.lifetime, log10(m), z);
     *tau = pow(10, *tau);
   }
 
@@ -171,7 +172,7 @@ PyObject* pymass_from_lifetime(PyObject* self, PyObject* args) {
   
   /* Init stellar physics */
   struct feedback_props fp;
-  feedback_props_init(&fp, &pconst, &hydro_props, &us, &params, &cosmo);
+  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);
 
   /* return object */
   PyArrayObject *mass = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(time), PyArray_DIMS(time), NPY_FLOAT);
@@ -182,7 +183,7 @@ PyObject* pymass_from_lifetime(PyObject* self, PyObject* args) {
     float *m = (float*) PyArray_GETPTR1(mass, i);
 
     
-    *m = stellar_evolution_get_log_mass_from_lifetime(&fp.stellar_model.imf, log10(tau), z);
+    *m = stellar_evolution_get_log_mass_from_lifetime(&fp.stellar_model.lifetime, log10(tau), z);
     *m = pow(10, *m);
   }
 

tests/run_tests.sh

+#!/bin/bash
+
+set -e
+
+
+for filename in ./*.py;
+do
+    python $filename
+done

tests/test_initial_mass_function.py

+#!/usr/bin/env python3
+
+from pyswiftsim import libstellar
+import pyswiftsim
+
+import numpy as np
+import matplotlib.pyplot as plt
+plt.rc('text', usetex=True)
+
+
+N = 100
+
+
+if __name__ == "__main__":
+
+    with pyswiftsim.ParameterFile() as filename:
+
+        parser = pyswiftsim.parseYamlFile(filename)
+        # imf = parser["GEARInitialMassFunction"]
+        mass_min = float(1.)
+        mass_max = float(50.)
+
+        mass = np.linspace(mass_min, mass_max, N, dtype=np.float32)
+
+        print("Computing initial mass function...")
+
+        imf = libstellar.initialMassFunction(filename, mass)
+
+        print("IMF obtained: {}".format(imf))