Segregation Theil index

theil_seg.py

import numpy as np
import multiprocessing as mp

def theil_seg(g):
    '''
    Compute segregation Theil index for one observation
    ...

    g       : array
              Array of shape (g, ) with counts for each group in the tract
    '''
    c_ig = g / g.sum()
    lc_ig = np.emath.logn(g.shape[0], 1./c_ig)
    return (c_ig * lc_ig).sum()

def theil_seg_several(ng):
    '''
    Compute segregation Theil index for several observations
    ...

    ng      : array
              Array of shape (n, g) where n is the number of observations
              (e.g. tracts) for the Theil to be computed on and g is the
              number of groups per tract
    '''
    pool = mp.Pool(mp.cpu_count())
    return pool.map(theil_seg, ng)