# Example: Millenium development goal 1; WHO childhood hunger data

import pandas as pd

hunger = pd.read_csv('http://apps.who.int/gho/athena/data/GHO/WHOSIS_000008.csv?profile=text&filter=COUNTRY:;SEX:')

hunger = hunger[hunger['Sex'] != 'Both sexes']
# the last entry is all NaN 
hunger = hunger[hunger['Year'].notnull()]

hunger.head()

# emulates abline function, only possible to set basic line styles and width
def abline(intercept, gradient, *args, **kwargs):
    a = gca()
    xlim = a.get_xlim()
    ylim = a.get_ylim()
    
    if args:
        sty = args[0]
    else:
        sty = 'r'
        
    if kwargs:
        lw = kwargs['linewidth']
    else:
        lw = 5

    a.plot(xlim, [intercept + gradient * x for x in xlim], sty, linewidth=lw)
    a.set_xlim(xlim)
    a.set_ylim(ylim);

from statsmodels.formula.api import ols

lm1 = ols('Numeric ~ Year', hunger).fit()

plot(hunger['Year'], hunger['Numeric'], 'ob', alpha=0.6)
plot(hunger['Year'], lm1.fittedvalues, 'grey', linewidth=3);

sex_groups = hunger.groupby('Sex')

c = {'Female' : 'r', 'Male' : 'k'}

idM = hunger['Sex'] == 'Male'
idF = hunger['Sex'] == 'Female'

lmM = ols('Numeric ~ Year', hunger[idM]).fit()
lmF = ols('Numeric ~ Year', hunger[idF]).fit()

for sex, df in sex_groups:
    scatter(df['Year'], df['Numeric'], c=c[sex], alpha=.6)

plot(hunger['Year'][idM], lmM.fittedvalues, c['Male'])
plot(hunger['Year'][idF], lmF.fittedvalues, c['Female']);

lmBoth = ols('Numeric ~ Year + Sex', hunger).fit()

for sex, df in sex_groups:
    scatter(df['Year'], df['Numeric'], c=c[sex], alpha=.6)

abline(lmBoth.params['Intercept'], 
       lmBoth.params['Year'], c['Female'], linewidth=3)
abline(lmBoth.params['Intercept'] + lmBoth.params['Sex[T.Male]'], 
       lmBoth.params['Year'], c['Male'], linewidth=3)

lmBoth = ols('Numeric ~ Year + Sex + Sex * Year', hunger).fit()

for sex, df in sex_groups:
    scatter(df['Year'], df['Numeric'], c=c[sex], alpha=.6)

abline(lmBoth.params['Intercept'], 
       lmBoth.params['Year'], c['Female'], linewidth=3)
abline(lmBoth.params['Intercept'] + lmBoth.params['Sex[T.Male]'], 
       lmBoth.params['Year'] + lmBoth.params['Sex[T.Male]:Year'], c['Male'], linewidth=3)

lmBoth.summary()