python lasso_missing_py2.py

Posted 2021-05-09

from pymc import *
import numpy as np
import pandas as pd
from numpy.ma import masked_values

# Import data, filling missing values with sentinels (-999)
test_scores = pd.read_csv('data/test_scores.csv')

# Extract variables: test score, gender, number of siblings, previous disability, age, 
# mother with HS education or better, hearing loss identified by 3 months of age
(score, male, siblings, disability, 
    age, mother_hs, early_ident) = test_scores[['score', 'male', 'sib', 
                                                'synd_or_disab', 'age_test',
                                                'mother_hs', 'ident_by_3']].astype(float).values.T

def generate_model():
    
    # Imputation of maternal education
    p_mother_hs = Beta("p_mother_hs", 1, 1, value=0.5)
    mother_hs_masked = masked_values(mother_hs, value=np.nan)
    print(mother_hs_masked)
    x_mother_hs = Bernoulli('x_mother_hs', p_mother_hs, value=mother_hs_masked, observed=True)
    print(x_mother_hs.value)
    
    # Imputation of previous disability
    p_synd_or_disab = Beta("p_synd_or_disab", 1, 1, value=0.5)
    synd_or_disab_masked = masked_values(disability, value=np.nan)
    x_synd_or_disab = Bernoulli('x_synd_or_disab', p_synd_or_disab, value=synd_or_disab_masked, observed=True)
    
    # Imputation of siblings
    theta_sib = Exponential("p_sib", 1, value=1)
    sib_masked = masked_values(siblings, value=np.nan)
    x_sib = Poisson('x_sib', theta_sib, value=sib_masked, observed=True)

    # Fixed effects
    beta = Laplace("beta", 0, 100, value=[0]*7)

    @deterministic
    def theta(b=beta, x_sib=x_sib, x_mother_hs=x_mother_hs, x_synd_or_disab=x_synd_or_disab):
        
        X = [male, x_sib, x_synd_or_disab, age, x_mother_hs, early_ident]
        return b[0] + np.dot(b[1:], X) 
    
    sigma = HalfCauchy("sigma", 0, 5, value=10)
    tau = sigma**-2
    
    # Data likelihood
    observed_score = Normal("observed_score", mu=theta, tau=tau, value=score, observed=True)
    
    return locals()


def run(n=20000):
    if n == 'short':
        n = 100
    M = MCMC(generate_model())
    M.sample(n, burn=n/2)
    Matplot.summary_plot(M.beta)


if __name__ == '__main__':
    run()