kroupaimf.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on Mon Apr 25 10:35:03 2016

@author: mints
"""
import numpy as np

class KroupaIMF(object):
    def __init__(self):
        self._bins = np.array([0.01, 0.08, 0.5, 150.])
        self._slopes = np.array([-0.3, -1.3, -2.3])
        self.norm = []
        self.fraction = np.ones(len(self._slopes))
        self._normalize()

    def _normalize(self):
        koeff = [1.]
        for i in range(1, len(self._bins) - 1):
            koeff.append(self._bins[i]**(self._slopes[i-1] - self._slopes[i]))
        koeff = np.cumprod(koeff)
        total_sum = 0.
        for i in range(len(self._slopes)):
            self.fraction[i] =  (self._bins[i + 1]**(self._slopes[i] + 1.) - \
                 self._bins[i]**(self._slopes[i] + 1.)) / (self._slopes[i] + 1.)
            total_sum = total_sum + koeff[i] * self.fraction[i]
        self.norm = koeff / total_sum
        self.fraction = self.norm*self.fraction

    def get_value(self, x, abin=None):
        if abin is None:
            abin = self._bins.searchsorted(x) - 1
        return x ** self._slopes[abin] * self.norm[abin]

    def get_value_int(self, x, abin=None):
        if abin is None:
            abin = self._bins.searchsorted(x) - 1
        return x ** (self._slopes[abin] + 1) * self.norm[abin]

    def integrate(self, low=None, up=None):
        if low is None or low < self._bins[0]:
            low = self._bins[0]
        if up is None or up > self._bins[-1]:
            up = self._bins[-1]
        result = 0
        low_bin = self._bins.searchsorted(low) - 1
        if low_bin < 0:
            low_bin = 0
        up_bin = self._bins.searchsorted(up) - 1
        #print low, up, low_bin, up_bin
        if low_bin == up_bin:
            return (self.get_value_int(up, low_bin) -
                    self.get_value_int(low, low_bin)) / (self._slopes[low_bin] + 1)
        else:
            result = (self.get_value_int(self._bins[low_bin + 1], low_bin) -
                      self.get_value_int(low, low_bin)) / (self._slopes[low_bin] + 1)
            #print result, low_bin
            for abin in range(low_bin + 1, up_bin):
                result = result + self.fraction[abin]
            result = result + (self.get_value_int(up, up_bin) -
                               self.get_value_int(self._bins[up_bin], up_bin)) / (self._slopes[up_bin] + 1)
            return result