stellar_evolution_wrapper.c

/*******************************************************************************
 * This file is part of PYSWIFTSIM.
 * 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/>.
 *
 ******************************************************************************/
#include "stellar_evolution_wrapper.h"
#include "pyswiftsim_tools.h"

/**
 * @brief Compute the initial mass function.
 *
 * args is expecting:
 *  - the name of the parameter file
 *  - a numpy array containing the masses.
 *
 * @param self calling object
 * @param args arguments
 * @return The initial mass function
 */
PyObject* pyinitial_mass_function(PyObject* self, PyObject* args) {
  import_array();
  

  char *filename;

  PyArrayObject *mass;

  /* parse argument */
  if (!PyArg_ParseTuple(
        args, "sO", &filename, &mass))
    return NULL;

  /* check numpy array */
  if (pytools_check_array(mass, 1, NPY_FLOAT, "Mass") != SWIFT_SUCCESS)
    return NULL;

  size_t N = PyArray_DIM(mass, 0);

  int with_cosmo = 0;
  PYSWIFTSIM_INIT_STRUCTS();

  /* Init stellar physics */
  struct feedback_props fp;
  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);

  for(size_t i = 0; i < N; i++) {
    float m = *(float*) PyArray_GETPTR1(mass, i);
    float *imf_i = (float*) PyArray_GETPTR1(imf, i);

    /* change units internal units -> solar mass */
    m /= pconst.const_solar_mass;

    if (m > fp.stellar_model.imf.mass_max) {
      *imf_i = 0.;
      continue;
    }

    *imf_i = initial_mass_function_get_imf(&fp.stellar_model.imf, m);

    /* change from internal units to solar mass */
    *imf_i /= pconst.const_solar_mass;
  }

  return (PyObject *) imf;
  
}

/**
 * @brief Compute the lifetime of a star.
 *
 * args is expecting:
 *  - the name of the parameter file
 *  - a numpy array containing the masses.
 *  - a numpy array containing the metallicity.
 *
 * @param self calling object
 * @param args arguments
 * @return the lifetime of the stars
 */
PyObject* pylifetime_from_mass(PyObject* self, PyObject* args) {
  import_array();
  
  char *filename;

  PyArrayObject *mass;
  PyArrayObject *metallicity;

  /* parse argument */
  if (!PyArg_ParseTuple(
      args, "sOO", &filename, &mass, &metallicity))
    return NULL;

  /* check numpy array */
  if (pytools_check_array(mass, 1, NPY_FLOAT, "Mass") != SWIFT_SUCCESS)
    return NULL;
  if (pytools_check_array(metallicity, 1, NPY_FLOAT, "Metallicity") != SWIFT_SUCCESS)
    return NULL;

  if (PyArray_DIM(metallicity, 0) != PyArray_DIM(mass, 0))
    error("Mass and metallicity should have the same dimension");


  size_t N = PyArray_DIM(mass, 0);

  int with_cosmo = 0;

  PYSWIFTSIM_INIT_STRUCTS();
  
  /* Init stellar physics */
  struct feedback_props fp;
  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);

  for(size_t i = 0; i < N; i++) {
    float m = *(float*) PyArray_GETPTR1(mass, i);
    float z = *(float*) PyArray_GETPTR1(metallicity, i);
    float *tau = (float*) PyArray_GETPTR1(lifetime, i);


    /* change units internal units -> solar mass */
    m /= pconst.const_solar_mass;
    
    *tau = lifetime_get_log_lifetime_from_mass(&fp.stellar_model.lifetime, log10(m), z);
    *tau = pow(10, *tau) * pconst.const_year * 1e6;
  }

  return (PyObject *) lifetime;
  
}

/**
 * @brief Compute the mass of a star with a given lifetime.
 *
 * args is expecting:
 *  - the name of the parameter file
 *  - a numpy array containing the time.
 *  - a numpy array containing the metallicity.
 *
 * @param self calling object
 * @param args arguments
 * @return The mass of the stars.
 */
PyObject* pymass_from_lifetime(PyObject* self, PyObject* args) {
  import_array();
  
  char *filename;

  PyArrayObject *time;
  PyArrayObject *metallicity;

  /* parse argument */
  if (!PyArg_ParseTuple(
      args, "sOO", &filename, &time, &metallicity))
    return NULL;

  /* check numpy array */
  if (pytools_check_array(time, 1, NPY_FLOAT, "Lifetime") != SWIFT_SUCCESS)
    return NULL;
  if (pytools_check_array(metallicity, 1, NPY_FLOAT, "Metallicity") != SWIFT_SUCCESS)
    return NULL;

  if (PyArray_DIM(metallicity, 0) != PyArray_DIM(time, 0))
    error("Lifetime and metallicity should have the same dimension");


  size_t N = PyArray_DIM(time, 0);

  int with_cosmo = 0;

  PYSWIFTSIM_INIT_STRUCTS();
  
  /* Init stellar physics */
  struct feedback_props fp;
  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);

  for(size_t i = 0; i < N; i++) {
    float tau = *(float*) PyArray_GETPTR1(time, i);
    float z = *(float*) PyArray_GETPTR1(metallicity, i);
    float *m = (float*) PyArray_GETPTR1(mass, i);

    
    /* change units internal units to Myr */    
    tau /= pconst.const_year * 1e6;

    *m = lifetime_get_log_mass_from_lifetime(&fp.stellar_model.lifetime, log10(tau), z);
    *m = pow(10, *m) * pconst.const_solar_mass;
  }

  return (PyObject *) mass;
  
}



/**
 * @brief Compute the supernovae Ia rate.
 *
 * args is expecting:
 *  - the name of the parameter file
 *  - a numpy array containing the lower limit for the time.
 *  - a numpy array containing the upper limit for the time.
 *  - a numpy array containing the metallicity.
 *
 * @param self calling object
 * @param args arguments
 * @return Rate of supernovae Ia
 */
PyObject* pysupernovae_ia_rate(PyObject* self, PyObject* args) {
  import_array();
  
  char *filename;

  PyArrayObject *time1;
  PyArrayObject *time2;
  PyArrayObject *metallicity;

  /* parse argument */
  if (!PyArg_ParseTuple(
      args, "sOOO", &filename, &time1, &time2, &metallicity))
    return NULL;

  /* check numpy array */
  if (pytools_check_array(time1, 1, NPY_FLOAT, "Time1") != SWIFT_SUCCESS)
    return NULL;
  if (pytools_check_array(time2, 1, NPY_FLOAT, "Time2") != SWIFT_SUCCESS)
    return NULL;
  if (pytools_check_array(metallicity, 1, NPY_FLOAT, "Metallicity") != SWIFT_SUCCESS)
    return NULL;

  if (PyArray_DIM(metallicity, 0) != PyArray_DIM(time1, 0))
    error("Time1 and metallicity should have the same dimension");
  if (PyArray_DIM(metallicity, 0) != PyArray_DIM(time2, 0))
    error("Time2 and metallicity should have the same dimension");


  size_t N = PyArray_DIM(time1, 0);

  int with_cosmo = 0;

  PYSWIFTSIM_INIT_STRUCTS();
  
  /* Init stellar physics */
  struct feedback_props fp;
  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);

  /* return object */
  PyArrayObject *rate = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(time1), PyArray_DIMS(time1), NPY_FLOAT);

  for(size_t i = 0; i < N; i++) {
    float t1 = *(float*) PyArray_GETPTR1(time1, i);
    float t2 = *(float*) PyArray_GETPTR1(time2, i);
    float z = *(float*) PyArray_GETPTR1(metallicity, i);
    float *r = (float*) PyArray_GETPTR1(rate, i);

    /* change units internal units -> Myr */
    t1 /= pconst.const_year * 1e6;
    t2 /= pconst.const_year * 1e6;
    
    float m1 = lifetime_get_log_mass_from_lifetime(&fp.stellar_model.lifetime, log10(t1), z);
    m1 = pow(10, m1);
    float m2 = lifetime_get_log_mass_from_lifetime(&fp.stellar_model.lifetime, log10(t2), z);
    m2 = pow(10, m2);

    /* Need to reverse 1 and 2 due to the lifetime being a decreasing function */
    *r = supernovae_ia_get_number_per_unit_mass(&fp.stellar_model.snia, m2, m1) / (t2 - t1);

    /* change units 1 / (solMass Myr) -> internal units */
    *r /= pconst.const_solar_mass * pconst.const_year * 1e6;
  }

  return (PyObject *) rate;
  
}

/**
 * @brief Compute the supernovae II rate.
 *
 * args is expecting:
 *  - the name of the parameter file
 *  - a numpy array containing the lower limit for the time.
 *  - a numpy array containing the upper limit for the time.
 *  - a numpy array containing the metallicity.
 *
 * @param self calling object
 * @param args arguments
 * @return Rate of supernovae II
 */
PyObject* pysupernovae_ii_rate(PyObject* self, PyObject* args) {
  import_array();
  
  char *filename;

  PyArrayObject *time1;
  PyArrayObject *time2;
  PyArrayObject *metallicity;

  /* parse argument */
  if (!PyArg_ParseTuple(
      args, "sOOO", &filename, &time1, &time2, &metallicity))
    return NULL;

  /* check numpy array */
  if (pytools_check_array(time1, 1, NPY_FLOAT, "Time1") != SWIFT_SUCCESS)
    return NULL;
  if (pytools_check_array(time2, 1, NPY_FLOAT, "Time2") != SWIFT_SUCCESS)
    return NULL;
  if (pytools_check_array(metallicity, 1, NPY_FLOAT, "Metallicity") != SWIFT_SUCCESS)
    return NULL;

  if (PyArray_DIM(metallicity, 0) != PyArray_DIM(time1, 0))
    error("Time1 and metallicity should have the same dimension");
  if (PyArray_DIM(metallicity, 0) != PyArray_DIM(time2, 0))
    error("Time2 and metallicity should have the same dimension");


  size_t N = PyArray_DIM(time1, 0);

  int with_cosmo = 0;

  PYSWIFTSIM_INIT_STRUCTS();
  
  /* Init stellar physics */
  struct feedback_props fp;
  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);

  /* return object */
  PyArrayObject *rate = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(time1), PyArray_DIMS(time1), NPY_FLOAT);

  for(size_t i = 0; i < N; i++) {
    float t1 = *(float*) PyArray_GETPTR1(time1, i);
    float t2 = *(float*) PyArray_GETPTR1(time2, i);
    float z = *(float*) PyArray_GETPTR1(metallicity, i);
    float *r = (float*) PyArray_GETPTR1(rate, i);

    /* change units internal units -> Myr */
    t1 /= pconst.const_year * 1e6;
    t2 /= pconst.const_year * 1e6;    
    
    float m1 = lifetime_get_log_mass_from_lifetime(&fp.stellar_model.lifetime, log10(t1), z);
    m1 = pow(10, m1);
    float m2 = lifetime_get_log_mass_from_lifetime(&fp.stellar_model.lifetime, log10(t2), z);
    m2 = pow(10, m2);

    /* Need to reverse 1 and 2 due to the lifetime being a decreasing function */
    *r = supernovae_ii_get_number_per_unit_mass(&fp.stellar_model.snii, m2, m1) / (t2 - t1);

    /* change units 1 / (solMass Myr) -> internal units */
    *r /= pconst.const_solar_mass * pconst.const_year * 1e6;

  }

  return (PyObject *) rate;
  
}

/**
 * @brief Compute the supernovae II yields.
 *
 * args is expecting:
 *  - the name of the parameter file
 *  - a numpy array containing the lower limit for the mass.
 *  - a numpy array containing the upper limit for the mass.
 *
 * @param self calling object
 * @param args arguments
 * @return mass fraction of ejected metals
 */
PyObject* pysupernovae_ii_yields(PyObject* self, PyObject* args) {
  import_array();
  
  char *filename;

  PyArrayObject *mass1;
  PyArrayObject *mass2;

  /* parse argument */
  if (!PyArg_ParseTuple(
      args, "sOO", &filename, &mass1, &mass2))
    return NULL;

  /* check numpy array */
  if (pytools_check_array(mass1, 1, NPY_FLOAT, "Mass1") != SWIFT_SUCCESS)
    return NULL;
  if (pytools_check_array(mass2, 1, NPY_FLOAT, "Mass2") != SWIFT_SUCCESS)
    return NULL;

  if (PyArray_NDIM(mass1) != 1 || PyArray_NDIM(mass2) != 1)
    error("The inputs should be 1-D arrays");

  if (PyArray_DIM(mass1, 0) != PyArray_DIM(mass2, 0))
    error("Mass1 and Mass2 should have the same dimension");


  size_t N = PyArray_DIM(mass1, 0);

  int with_cosmo = 0;

  PYSWIFTSIM_INIT_STRUCTS();
  
  /* Init stellar physics */
  struct feedback_props fp;
  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);

  /* return object */
  const int ndims = 2;
  npy_intp dims[2] = {N, GEAR_CHEMISTRY_ELEMENT_COUNT};
  PyArrayObject *yields_out = (PyArrayObject *)PyArray_SimpleNew(ndims, dims, NPY_FLOAT);

  for(size_t i = 0; i < N; i++) {
    float yields[GEAR_CHEMISTRY_ELEMENT_COUNT] = {0.};
    const float m1 = *(float*) PyArray_GETPTR1(mass1, i) / pconst.const_solar_mass;
    const float log_m1 = log10(m1);
    const float m2 = *(float*) PyArray_GETPTR1(mass2, i) / pconst.const_solar_mass;
    const float log_m2 = log10(m2);

    /* Compute the yields */
    supernovae_ii_get_yields_from_integral(&fp.stellar_model.snii, log_m1, log_m2, yields);

    for(int j = 0; j < GEAR_CHEMISTRY_ELEMENT_COUNT; j++) {
      float *y = (float*) PyArray_GETPTR2(yields_out, i, j);
      *y = yields[j];
    }
  }

  return (PyObject *) yields_out;
  
}

/**
 * @brief Compute the mass ejected by a supernovae II.
 *
 * args is expecting:
 *  - the name of the parameter file
 *  - a numpy array containing the lower limit for the mass.
 *  - a numpy array containing the upper limit for the mass.
 *
 * @param self calling object
 * @param args arguments
 * @return mass fraction of ejected mass.
 */
PyObject* pysupernovae_ii_mass_ejected(PyObject* self, PyObject* args) {
  import_array();
  
  char *filename;

  PyArrayObject *mass1;
  PyArrayObject *mass2;

  /* parse argument */
  if (!PyArg_ParseTuple(
      args, "sOO", &filename, &mass1, &mass2))
    return NULL;

  /* check numpy array */
  if (pytools_check_array(mass1, 1, NPY_FLOAT, "Mass1") != SWIFT_SUCCESS)
    return NULL;
  if (pytools_check_array(mass2, 1, NPY_FLOAT, "Mass2") != SWIFT_SUCCESS)
    return NULL;

  if (PyArray_NDIM(mass1) != 1 || PyArray_NDIM(mass2) != 1)
    error("The inputs should be 1-D arrays");

  if (PyArray_DIM(mass1, 0) != PyArray_DIM(mass2, 0))
    error("Mass1 and Mass2 should have the same dimension");


  size_t N = PyArray_DIM(mass1, 0);

  int with_cosmo = 0;

  PYSWIFTSIM_INIT_STRUCTS();
  
  /* Init stellar physics */
  struct feedback_props fp;
  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);

  /* return object */
  PyArrayObject *mass_ejected = (PyArrayObject *)PyArray_SimpleNew(
      PyArray_NDIM(mass1), PyArray_DIMS(mass1), NPY_FLOAT);

  for(size_t i = 0; i < N; i++) {
    const float m1 = *(float*) PyArray_GETPTR1(mass1, i) / pconst.const_solar_mass;
    const float log_m1 = log10(m1);
    const float m2 = *(float*) PyArray_GETPTR1(mass2, i) / pconst.const_solar_mass;
    const float log_m2 = log10(m2);
    float *y = (float*) PyArray_GETPTR1(mass_ejected, i);

    *y = supernovae_ii_get_ejected_mass_fraction_processed_from_integral(&fp.stellar_model.snii, log_m1, log_m2);

  }

  return (PyObject *) mass_ejected;
  
}

/**
 * @brief Compute the non processed mass ejected by a supernovae II.
 *
 * args is expecting:
 *  - the name of the parameter file
 *  - a numpy array containing the lower limit for the mass.
 *  - a numpy array containing the upper limit for the mass.
 *
 * @param self calling object
 * @param args arguments
 * @return mass ratio of (non processed) mass ejected
 */
PyObject* pysupernovae_ii_mass_ejected_non_processed(PyObject* self, PyObject* args) {
  import_array();
  
  char *filename;

  PyArrayObject *mass1;
  PyArrayObject *mass2;

  /* parse argument */
  if (!PyArg_ParseTuple(
      args, "sOO", &filename, &mass1, &mass2))
    return NULL;

  /* check numpy array */
  if (pytools_check_array(mass1, 1, NPY_FLOAT, "Mass1") != SWIFT_SUCCESS)
    return NULL;
  if (pytools_check_array(mass2, 1, NPY_FLOAT, "Mass2") != SWIFT_SUCCESS)
    return NULL;

  if (PyArray_NDIM(mass1) != 1 || PyArray_NDIM(mass2) != 1)
    error("The inputs should be 1-D arrays");

  if (PyArray_DIM(mass1, 0) != PyArray_DIM(mass2, 0))
    error("Mass1 and Mass2 should have the same dimension");


  size_t N = PyArray_DIM(mass1, 0);

  int with_cosmo = 0;

  PYSWIFTSIM_INIT_STRUCTS();
  
  /* Init stellar physics */
  struct feedback_props fp;
  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);

  /* return object */
  PyArrayObject *mass_ejected = (PyArrayObject *)PyArray_SimpleNew(
      PyArray_NDIM(mass1), PyArray_DIMS(mass1), NPY_FLOAT);

  for(size_t i = 0; i < N; i++) {
    const float m1 = *(float*) PyArray_GETPTR1(mass1, i) / pconst.const_solar_mass;
    const float log_m1 = log10(m1);
    const float m2 = *(float*) PyArray_GETPTR1(mass2, i) / pconst.const_solar_mass;
    const float log_m2 = log10(m2);
    float *y = (float*) PyArray_GETPTR1(mass_ejected, i);

    *y = supernovae_ii_get_ejected_mass_fraction_non_processed_from_integral(&fp.stellar_model.snii, log_m1, log_m2);
    
  }

  return (PyObject *) mass_ejected;
  
}

/**
 * @brief Compute the supernovae Ia yields.
 *
 * args is expecting:
 *  - the name of the parameter file
 *
 * @param self calling object
 * @param args arguments
 * @return Mass fraction of the supernovae Ia yields.
 */
PyObject* pysupernovae_ia_yields(PyObject* self, PyObject* args) {
  import_array();
  
  char *filename;

  /* parse argument */
  if (!PyArg_ParseTuple(
      args, "s", &filename))
    return NULL;

  int with_cosmo = 0;

  PYSWIFTSIM_INIT_STRUCTS();
  
  /* Init stellar physics */
  struct feedback_props fp;
  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);

  /* return object */
  const int ndims = 1;
  npy_intp dims[1] = {GEAR_CHEMISTRY_ELEMENT_COUNT};
  PyArrayObject *yields_out = (PyArrayObject *)PyArray_SimpleNew(ndims, dims, NPY_FLOAT);

  const float *yields = supernovae_ia_get_yields(&fp.stellar_model.snia);

  for(size_t i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
    float *y = (float*) PyArray_GETPTR1(yields_out, i);

    *y = yields[i] / fp.stellar_model.snia.mass_white_dwarf;
  }

  return (PyObject *) yields_out;
  
}

/**
 * @brief Compute the mass ejected by a supernovae Ia.
 *
 * args is expecting:
 *  - the name of the parameter file
 *
 * @param self calling object
 * @param args arguments
 * @return mass fraction ejected.
 */
PyObject* pysupernovae_ia_mass_ejected(PyObject* self, PyObject* args) {
  import_array();
  
  char *filename;

  /* parse argument */
  if (!PyArg_ParseTuple(
      args, "s", &filename))
    return NULL;

  int with_cosmo = 0;

  PYSWIFTSIM_INIT_STRUCTS();
  
  /* Init stellar physics */
  struct feedback_props fp;
  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);

  /* return object */
  const int ndims = 1;
  npy_intp dims[1] = {1};
  PyArrayObject *mass_ejected = (PyArrayObject *)PyArray_SimpleNew(ndims, dims, NPY_FLOAT);

  float *y = (float*) PyArray_GETPTR1(mass_ejected, 0);

  *y = supernovae_ia_get_ejected_mass_processed(&fp.stellar_model.snia) / fp.stellar_model.snia.mass_white_dwarf;

  return (PyObject *) mass_ejected;
  
}


/**
 * @brief Read the elements names.
 *
 * args is expecting:
 *  - the name of the parameter file
 *
 * @param self calling object
 * @param args arguments
 * @return cooling rate
 */
PyObject* pystellar_get_elements(PyObject* self, PyObject* args) {
  import_array();
  
  char *filename;

  /* parse argument */
  if (!PyArg_ParseTuple(args, "s", &filename))
    return NULL;

  int with_cosmo = 0;

  PYSWIFTSIM_INIT_STRUCTS();
  
  /* Init stellar physics */
  struct feedback_props fp;
  feedback_props_init(&fp, &pconst, &us, &params, &hydro_props, &cosmo);

  /* Save the variables */
  PyObject *t = PyList_New(GEAR_CHEMISTRY_ELEMENT_COUNT);

  for(int i = 0; i < GEAR_CHEMISTRY_ELEMENT_COUNT; i++) {
    const char *name = stellar_evolution_get_element_name(&fp.stellar_model, i);
    PyObject *string = PyUnicode_FromString(name);
    int test = PyList_SetItem(t, i, string);
    if (test != 0) {
      error("Failed to add an element");
    }
  }

  return t;
  
}