<s> """<STR_LIT>""" <EOL> from __future__ import print_function <EOL> import os <EOL> import time <EOL> import gzip <EOL> from zipfile import ZipFile <EOL> from urllib import urlretrieve <EOL> import cPickle as pickle <EOL> import operator <EOL> from collections import namedtuple , defaultdict <EOL> import numpy as np <EOL> from functools import reduce <EOL> np . random . seed ( <NUM_LIT:0> ) <EOL> from scipy . special import gammaln <EOL> from scipy . misc import logsumexp <EOL> import matplotlib . pyplot as plt <EOL> from hips . plotting . layout import create_figure , create_axis_at_location <EOL> from hips . plotting . sausage import sausage_plot <EOL> import brewer2mpl <EOL> colors = brewer2mpl . get_map ( "<STR_LIT>" , "<STR_LIT>" , <NUM_LIT:9> ) . mpl_colors <EOL> goodcolors = np . array ( [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:4> , <NUM_LIT:6> , <NUM_LIT:7> , <NUM_LIT:8> ] ) <EOL> colors = np . array ( colors ) [ goodcolors ] <EOL> from GPy . kern import RBF , Matern52 <EOL> from pgmult . utils import pi_to_psi , psi_to_pi , ln_pi_to_psi <EOL> import pgmult . gp <EOL> reload ( pgmult . gp ) <EOL> import pgmult . distributions <EOL> reload ( pgmult . distributions ) <EOL> def load_data ( <EOL> start_train_year = <NUM_LIT> , end_train_year = <NUM_LIT> , end_year = <NUM_LIT> , <EOL> continental = False , DC = True , N_names = None , train_state = None , <EOL> downsample = None ) : <EOL> data , years , states , names = download_data ( N_names , end_year ) <EOL> loc_file = os . path . join ( '<STR_LIT:data>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> with open ( loc_file ) as infile : <EOL> latlon = { state : ( float ( lat ) , float ( lon ) ) for <EOL> state , lat , lon in <EOL> [ line . strip ( ) . split ( '<STR_LIT:U+002C>' ) for line in infile ] } <EOL> flat_data , flat_years , flat_states , lat , lon = all_flat = flatten ( data , years , states , names , latlon ) <EOL> train_inds , test_inds = get_train_test_ind ( <EOL> flat_years , flat_states , train_state , continental , DC ) <EOL> train_data , train_years , train_states , train_lat , train_lon = map ( lambda arr : arr [ train_inds ] , all_flat ) <EOL> test_data , test_years , test_states , test_lat , test_lon = map ( lambda arr : arr [ test_inds ] , all_flat ) <EOL> if downsample is not None : <EOL> print ( "<STR_LIT>" , downsample , "<STR_LIT>" ) <EOL> assert isinstance ( downsample , int ) and downsample > <NUM_LIT:0> <EOL> from pgmult . utils import downsample_data <EOL> downsample_train_data = downsample_data ( train_data , downsample ) <EOL> downsample_test_data = downsample_data ( test_data , downsample ) <EOL> else : <EOL> downsample_train_data = downsample_test_data = None <EOL> Dataset = namedtuple ( "<STR_LIT>" , [ "<STR_LIT>" , "<STR_LIT:data>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] ) <EOL> train = Dataset ( N_names , train_data , train_years , train_lat , train_lon , train_states , names ) <EOL> test = Dataset ( N_names , test_data , test_years , test_lat , test_lon , test_states , names ) <EOL> downsample_train = Dataset ( <EOL> N_names , downsample_train_data , train_years , train_lat , train_lon , train_states , names ) <EOL> downsample_test = Dataset ( <EOL> N_names , downsample_test_data , test_years , test_lat , test_lon , test_states , names ) <EOL> return downsample_train , downsample_test , train , test <EOL> def download_data ( N_names , end_year ) : <EOL> url = '<STR_LIT>' <EOL> datafile = os . path . join ( "<STR_LIT:data>" , "<STR_LIT>" , "<STR_LIT>" ) <EOL> if not os . path . exists ( datafile ) : <EOL> print ( '<STR_LIT>' ) <EOL> urlretrieve ( url , datafile ) <EOL> print ( '<STR_LIT>' ) <EOL> print ( '<STR_LIT>' ) <EOL> alldata = parse_names_files ( ZipFile ( datafile ) , N_names , end_year ) <EOL> print ( '<STR_LIT>' ) <EOL> return alldata <EOL> def parse_names_files ( zfile , N_names , end_year ) : <EOL> def parse_state ( string ) : <EOL> data = defaultdict ( lambda : defaultdict ( int ) ) <EOL> for line in string . strip ( ) . split ( '<STR_LIT:\r\n>' ) : <EOL> state , gender , year , name , count = line . split ( '<STR_LIT:U+002C>' ) <EOL> if gender . lower ( ) == '<STR_LIT:m>' and int ( year ) <= end_year : <EOL> data [ name . lower ( ) ] [ int ( year ) ] += int ( count ) <EOL> return data <EOL> def get_years ( dct ) : <EOL> sum = lambda lst : reduce ( operator . or_ , lst ) <EOL> return sorted ( sum ( set ( dd . keys ( ) ) for d in dct . values ( ) for dd in d . values ( ) ) ) <EOL> def get_top_names ( dct , N_names ) : <EOL> counts = defaultdict ( int ) <EOL> for statedict in dct . values ( ) : <EOL> for name , yeardict in statedict . iteritems ( ) : <EOL> counts [ name ] += sum ( yeardict . values ( ) ) <EOL> return sorted ( counts . keys ( ) , key = counts . __getitem__ , reverse = True ) [ : N_names ] <EOL> def get_data_array ( dct , years , states , names ) : <EOL> counts = np . zeros ( ( len ( years ) , len ( states ) , len ( names ) ) ) <EOL> for year_idx , year in enumerate ( years ) : <EOL> for state_idx , state in enumerate ( states ) : <EOL> for name_idx , name in enumerate ( names ) : <EOL> counts [ year_idx , state_idx , name_idx ] = dct [ state ] [ name ] [ year ] <EOL> return counts <EOL> keys = [ key for key in zfile . namelist ( ) if key . endswith ( '<STR_LIT>' ) ] <EOL> dct = { key [ : <NUM_LIT:2> ] : parse_state ( zfile . read ( key ) ) for key in keys } <EOL> years = get_years ( dct ) <EOL> states = sorted ( dct . keys ( ) ) <EOL> names = get_top_names ( dct , N_names ) <EOL> data = get_data_array ( dct , years , states , names ) <EOL> return data , years , states , names <EOL> def flatten ( data , years , states , names , latlon ) : <EOL> flat_data = data . reshape ( - <NUM_LIT:1> , len ( names ) ) <EOL> flat_years = np . repeat ( years , len ( states ) ) <EOL> flat_states = np . tile ( states , len ( years ) ) <EOL> lat , lon = np . array ( [ latlon [ st ] for st in flat_states ] ) . T <EOL> return flat_data , flat_years , flat_states , lat , lon <EOL> def get_train_test_ind ( flat_years , flat_states , train_state , continental , DC ) : <EOL> train_inds = ( flat_years >= start_train_year ) & ( flat_years < end_train_year ) <EOL> test_inds = ( flat_years >= end_train_year ) <EOL> if train_state is not None : <EOL> train_inds &= ( flat_states == train_state ) <EOL> test_inds &= ( flat_states == train_state ) <EOL> if continental : <EOL> not_AK = np . array ( [ st . upper ( ) != "<STR_LIT>" for st in flat_states ] ) <EOL> train_inds &= not_AK <EOL> test_inds &= not_AK <EOL> not_HI = np . array ( [ st . upper ( ) != "<STR_LIT>" for st in flat_states ] ) <EOL> train_inds &= not_HI <EOL> test_inds &= not_HI <EOL> if not DC : <EOL> not_DC = np . array ( [ st . upper ( ) != "<STR_LIT>" for st in flat_states ] ) <EOL> train_inds &= not_DC <EOL> test_inds &= not_DC <EOL> return train_inds , test_inds <EOL> def get_inputs ( data ) : <EOL> return np . hstack ( ( data . years [ : , None ] , data . lon [ : , None ] , data . lat [ : , None ] ) ) <EOL> def fit_gp_multinomial_model ( model , test , pi_train = None , N_samples = <NUM_LIT:100> , run = <NUM_LIT:1> ) : <EOL> if pi_train is not None : <EOL> if isinstance ( model , pgmult . gp . LogisticNormalGP ) : <EOL> model . data_list [ <NUM_LIT:0> ] [ "<STR_LIT>" ] = ln_pi_to_psi ( pi_train ) - model . mu <EOL> elif isinstance ( model , pgmult . gp . MultinomialGP ) : <EOL> model . data_list [ <NUM_LIT:0> ] [ "<STR_LIT>" ] = pi_to_psi ( pi_train ) - model . mu <EOL> model . resample_omega ( ) <EOL> else : <EOL> model . initialize_from_data ( ) <EOL> results_base = os . path . join ( "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" % run , "<STR_LIT>" ) <EOL> results_file = results_base + "<STR_LIT>" <EOL> if os . path . exists ( results_file ) : <EOL> with gzip . open ( results_file , "<STR_LIT:r>" ) as f : <EOL> samples , lls , pred_lls , timestamps = pickle . load ( f ) <EOL> else : <EOL> Z_test = get_inputs ( test ) <EOL> lls = [ model . log_likelihood ( ) ] <EOL> samples = [ model . copy_sample ( ) ] <EOL> pred_ll , pred_pi = model . predictive_log_likelihood ( Z_test , test . data ) <EOL> pred_lls = [ pred_ll ] <EOL> pred_pis = [ pred_pi ] <EOL> times = [ <NUM_LIT:0> ] <EOL> print ( "<STR_LIT>" , lls [ <NUM_LIT:0> ] ) <EOL> print ( "<STR_LIT>" , pred_lls [ <NUM_LIT:0> ] ) <EOL> for itr in xrange ( N_samples ) : <EOL> print ( "<STR_LIT>" , itr ) <EOL> tic = time . time ( ) <EOL> model . resample_model ( verbose = True ) <EOL> times . append ( time . time ( ) - tic ) <EOL> samples . append ( model . copy_sample ( ) ) <EOL> lls . append ( model . log_likelihood ( ) ) <EOL> pred_ll , pred_pi = model . predictive_log_likelihood ( get_inputs ( test ) , test . data ) <EOL> pred_lls . append ( pred_ll ) <EOL> pred_pis . append ( pred_pi ) <EOL> print ( "<STR_LIT>" , lls [ - <NUM_LIT:1> ] ) <EOL> print ( "<STR_LIT>" , pred_ll ) <EOL> lls = np . array ( lls ) <EOL> pred_lls = np . array ( pred_lls ) <EOL> timestamps = np . cumsum ( times ) <EOL> return samples , lls , pred_lls , pred_pis , timestamps <EOL> def compute_pred_likelihood ( model , samples , test ) : <EOL> Z_pred = get_inputs ( test ) <EOL> preds = [ ] <EOL> for sample in samples : <EOL> model . set_sample ( sample ) <EOL> preds . append ( model . predict ( Z_pred , full_output = True ) [ <NUM_LIT:1> ] ) <EOL> psi_pred_mean = np . mean ( preds , axis = <NUM_LIT:0> ) <EOL> if isinstance ( model , pgmult . gp . MultinomialGP ) : <EOL> pi_pred_mean = np . array ( [ psi_to_pi ( psi ) for psi in psi_pred_mean ] ) <EOL> elif isinstance ( model , pgmult . gp . LogisticNormalGP ) : <EOL> from pgmult . internals . utils import ln_psi_to_pi <EOL> pi_pred_mean = np . array ( [ ln_psi_to_pi ( psi ) for psi in psi_pred_mean ] ) <EOL> else : <EOL> raise NotImplementedError <EOL> pll_gp = gammaln ( test . data . sum ( axis = <NUM_LIT:1> ) + <NUM_LIT:1> ) . sum ( ) - gammaln ( test . data + <NUM_LIT:1> ) . sum ( ) <EOL> pll_gp += np . nansum ( test . data * np . log ( pi_pred_mean ) ) <EOL> return pll_gp <EOL> def compute_static_pred_ll ( static_model , test ) : <EOL> pll = <NUM_LIT:0> <EOL> years = np . unique ( test . years ) <EOL> X_test = test . data <EOL> for year in years : <EOL> pll += static_model . predictive_log_likelihood ( X_test [ test . years == year ] ) <EOL> return pll <EOL> def compute_fraction_top_names ( test , test_pis , top = <NUM_LIT:10> , bottom = <NUM_LIT:10> ) : <EOL> test_pi_mean = np . mean ( test_pis , axis = <NUM_LIT:0> ) <EOL> pred_perm = np . argsort ( - test_pi_mean , axis = <NUM_LIT:1> ) <EOL> data_perm = np . argsort ( - test . data , axis = <NUM_LIT:1> ) <EOL> top_scores = np . array ( [ len ( np . intersect1d ( pp [ : top ] , dp [ : top ] ) ) <EOL> for pp , dp in zip ( pred_perm , data_perm ) ] ) <EOL> bottom_scores = np . array ( [ len ( np . intersect1d ( pp [ - bottom : ] , dp [ - bottom : ] ) ) <EOL> for pp , dp in zip ( pred_perm , data_perm ) ] ) <EOL> results = np . zeros ( ( len ( np . unique ( test . years ) ) , <NUM_LIT:4> ) ) <EOL> for i , yr in enumerate ( sorted ( np . unique ( test . years ) ) ) : <EOL> results [ i , <NUM_LIT:0> ] = top_scores [ test . years == yr ] . mean ( <NUM_LIT:0> ) <EOL> results [ i , <NUM_LIT:1> ] = top_scores [ test . years == yr ] . std ( <NUM_LIT:0> ) <EOL> results [ i , <NUM_LIT:2> ] = bottom_scores [ test . years == yr ] . mean ( <NUM_LIT:0> ) <EOL> results [ i , <NUM_LIT:3> ] = bottom_scores [ test . years == yr ] . std ( <NUM_LIT:0> ) <EOL> print ( "<STR_LIT>" % ( yr , results [ i , <NUM_LIT:0> ] , results [ i , <NUM_LIT:1> ] ) ) <EOL> print ( "<STR_LIT>" % ( yr , results [ i , <NUM_LIT:2> ] , results [ i , <NUM_LIT:3> ] ) ) <EOL> return results <EOL> def plot_census_results ( train , samples , test , test_pis ) : <EOL> train_mus = np . array ( [ s [ <NUM_LIT:0> ] for s in samples ] ) <EOL> train_psis = np . array ( [ s [ <NUM_LIT:1> ] [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] for s in samples ] ) <EOL> train_psis += train_mus [ <NUM_LIT:0> ] [ None , None , : ] <EOL> train_pis = np . array ( [ psi_to_pi ( psi_sample ) for psi_sample in train_psis ] ) <EOL> train_pi_mean = np . mean ( train_pis , axis = <NUM_LIT:0> ) <EOL> train_pi_std = np . std ( train_pis , axis = <NUM_LIT:0> ) <EOL> test_pi_mean = np . mean ( test_pis , axis = <NUM_LIT:0> ) <EOL> test_pi_std = np . std ( test_pis , axis = <NUM_LIT:0> ) <EOL> train_pi_emp = train . data / train . data . sum ( axis = <NUM_LIT:1> ) [ : , None ] <EOL> test_pi_emp = test . data / test . data . sum ( axis = <NUM_LIT:1> ) [ : , None ] <EOL> names = [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] <EOL> states = [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] <EOL> linestyles = [ "<STR_LIT:->" , "<STR_LIT>" , "<STR_LIT::>" ] <EOL> fig = create_figure ( figsize = ( <NUM_LIT> , <NUM_LIT:3> ) ) <EOL> ax1 = create_axis_at_location ( fig , <NUM_LIT> , <NUM_LIT:0.5> , <NUM_LIT> , <NUM_LIT> ) <EOL> for name , color in zip ( names , colors ) : <EOL> for state , linestyle in zip ( states , linestyles ) : <EOL> train_state_inds = ( train . states == state ) <EOL> train_name_ind = np . array ( train . names ) == name . lower ( ) <EOL> train_years = train . years [ train . states == state ] <EOL> train_mean_name = train_pi_mean [ train_state_inds , train_name_ind ] <EOL> train_std_name = train_pi_std [ train_state_inds , train_name_ind ] <EOL> test_state_inds = ( test . states == state ) <EOL> test_name_ind = np . array ( test . names ) == name . lower ( ) <EOL> test_years = test . years [ test . states == state ] <EOL> test_mean_name = test_pi_mean [ test_state_inds , test_name_ind ] <EOL> test_std_name = test_pi_std [ test_state_inds , test_name_ind ] <EOL> years = np . concatenate ( ( train_years , test_years ) ) <EOL> mean_name = np . concatenate ( ( train_mean_name , test_mean_name ) ) <EOL> std_name = np . concatenate ( ( train_std_name , test_std_name ) ) <EOL> sausage_plot ( years , mean_name , std_name , <EOL> color = color , alpha = <NUM_LIT:0.5> ) <EOL> plt . plot ( years , mean_name , <EOL> color = color , label = "<STR_LIT>" % ( name , state ) , <EOL> ls = linestyle , lw = <NUM_LIT:2> ) <EOL> plt . plot ( train . years [ train_state_inds ] , <EOL> train_pi_emp [ train_state_inds , train_name_ind ] , <EOL> color = color , <EOL> ls = "<STR_LIT>" , marker = "<STR_LIT:x>" , markersize = <NUM_LIT:4> ) <EOL> plt . plot ( test . years [ test_state_inds ] , <EOL> test_pi_emp [ test_state_inds , test_name_ind ] , <EOL> color = color , <EOL> ls = "<STR_LIT>" , marker = "<STR_LIT:x>" , markersize = <NUM_LIT:4> ) <EOL> ylim = plt . gca ( ) . get_ylim ( ) <EOL> plt . plot ( ( test . years . min ( ) - <NUM_LIT:0.5> ) * np . ones ( <NUM_LIT:2> ) , ylim , '<STR_LIT>' , lw = <NUM_LIT:0.5> ) <EOL> plt . ylim ( ylim ) <EOL> plt . legend ( bbox_to_anchor = ( <NUM_LIT:0.> , <NUM_LIT> , <NUM_LIT:1.> , <NUM_LIT> ) , loc = <NUM_LIT:3> , <EOL> ncol = len ( names ) , mode = "<STR_LIT>" , borderaxespad = <NUM_LIT:0.> , <EOL> fontsize = "<STR_LIT>" ) <EOL> plt . xlabel ( "<STR_LIT>" ) <EOL> plt . xlim ( train . years . min ( ) , test . years . max ( ) + <NUM_LIT:0.1> ) <EOL> plt . ylabel ( "<STR_LIT>" ) <EOL> fig . savefig ( "<STR_LIT>" ) <EOL> plt . show ( ) <EOL> plt . pause ( <NUM_LIT:0.1> ) <EOL> def plot_spatial_distribution ( train , samples , name = "<STR_LIT>" , year = <NUM_LIT> ) : <EOL> mus = np . array ( [ s [ <NUM_LIT:0> ] for s in samples ] ) <EOL> psis = np . array ( [ s [ <NUM_LIT:1> ] [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] for s in samples ] ) <EOL> psis += mus [ <NUM_LIT:0> ] [ None , None , : ] <EOL> pis = np . array ( [ psi_to_pi ( psi_sample ) for psi_sample in psis ] ) <EOL> data = pis [ - <NUM_LIT:1> , train . years == year , train . names == name . lower ( ) ] <EOL> lons = train . lon [ train . years == year ] <EOL> lats = train . lat [ train . years == year ] <EOL> fig = plt . figure ( figsize = ( <NUM_LIT:3> , <NUM_LIT:3> ) ) <EOL> ax = fig . add_subplot ( <NUM_LIT> , aspect = "<STR_LIT>" ) <EOL> from mpl_toolkits . basemap import Basemap <EOL> m = Basemap ( width = <NUM_LIT> , height = <NUM_LIT> , <EOL> resolution = '<STR_LIT:l>' , projection = '<STR_LIT>' , <EOL> lat_ts = <NUM_LIT:50> , lat_0 = <NUM_LIT> , lon_0 = - <NUM_LIT> , <EOL> ax = ax ) <EOL> land_color = [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] <EOL> water_color = [ <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] <EOL> m . fillcontinents ( color = land_color , lake_color = water_color ) <EOL> m . drawcoastlines ( ) <EOL> m . drawstates ( ) <EOL> m . drawcountries ( ) <EOL> m . drawmapboundary ( fill_color = water_color ) <EOL> dx , dy = m ( lons , lats ) <EOL> glons , glats = m . makegrid ( <NUM_LIT:100> , <NUM_LIT:100> ) <EOL> gx , gy = m ( glons , glats ) <EOL> M = gx . size <EOL> from scipy . interpolate import griddata <EOL> gdata = griddata ( np . hstack ( ( dx [ : , None ] , dy [ : , None ] ) ) , <EOL> data , <EOL> np . hstack ( ( gx . reshape ( ( M , <NUM_LIT:1> ) ) , gy . reshape ( ( M , <NUM_LIT:1> ) ) ) ) , <EOL> method = "<STR_LIT>" ) <EOL> gdata = gdata . reshape ( gx . shape ) <EOL> cs = ax . contour ( gx , gy , gdata , <NUM_LIT:15> , cmap = "<STR_LIT>" , linewidth = <NUM_LIT:2> ) <EOL> plt . title ( "<STR_LIT>" % ( name , year ) ) <EOL> from mpl_toolkits . axes_grid1 import make_axes_locatable <EOL> divider = make_axes_locatable ( ax ) <EOL> cax = divider . append_axes ( "<STR_LIT:right>" , size = "<STR_LIT>" , pad = <NUM_LIT> ) <EOL> cbar = plt . colorbar ( cs , cax = cax ) <EOL> cbar . set_label ( "<STR_LIT>" , labelpad = <NUM_LIT:10> ) <EOL> plt . subplots_adjust ( left = <NUM_LIT> , bottom = <NUM_LIT:0.1> , top = <NUM_LIT> , right = <NUM_LIT> ) <EOL> fig . savefig ( "<STR_LIT>" % ( name . lower ( ) , year ) ) <EOL> return fig , ax , m <EOL> def plot_pred_ll_vs_time ( static_pll , emp_pll , <EOL> sbm_gp_timestamps , sbm_gp_plls , sbm_avg_pll , <EOL> lnm_gp_timestamps , lnm_gp_plls , lnm_avg_pll , <EOL> results_dir = "<STR_LIT:.>" , <EOL> begin_test_year = <NUM_LIT> , end_test_year = <NUM_LIT> , <EOL> norm = <NUM_LIT:1> , baseline = <NUM_LIT:0> , burnin = <NUM_LIT:0> <EOL> ) : <EOL> fig = plt . figure ( figsize = ( <NUM_LIT> * <NUM_LIT> , <NUM_LIT> * <NUM_LIT> ) ) <EOL> fig . set_tight_layout ( True ) <EOL> min_time = lnm_gp_timestamps [ burnin ] <EOL> max_time = sbm_gp_timestamps [ - <NUM_LIT:1> ] <EOL> plt . plot ( sbm_gp_timestamps [ burnin : ] , ( sbm_gp_plls [ burnin : ] - baseline ) / norm , lw = <NUM_LIT:2> , color = colors [ <NUM_LIT:0> ] , label = "<STR_LIT>" ) <EOL> plt . plot ( [ min_time , max_time ] , ( sbm_avg_pll - baseline ) / norm * np . ones ( <NUM_LIT:2> ) , ls = '<STR_LIT>' , color = colors [ <NUM_LIT:0> ] ) <EOL> plt . plot ( lnm_gp_timestamps [ burnin : ] , ( lnm_gp_plls [ burnin : ] - baseline ) / norm , lw = <NUM_LIT:2> , color = colors [ <NUM_LIT:1> ] , label = "<STR_LIT>" ) <EOL> plt . plot ( [ min_time , max_time ] , ( lnm_avg_pll - baseline ) / norm * np . ones ( <NUM_LIT:2> ) , ls = '<STR_LIT>' , color = colors [ <NUM_LIT:1> ] ) <EOL> plt . plot ( [ <NUM_LIT:0> , max_time ] , ( emp_pll - baseline ) / norm * np . ones ( <NUM_LIT:2> ) , lw = <NUM_LIT:2> , color = colors [ <NUM_LIT:2> ] , label = "<STR_LIT>" ) <EOL> plt . xscale ( "<STR_LIT>" ) <EOL> plt . xlim ( min_time , max_time ) <EOL> plt . legend ( loc = "<STR_LIT>" , prop = { "<STR_LIT:size>" : <NUM_LIT:8> } ) <EOL> plt . xlabel ( "<STR_LIT>" ) <EOL> plt . xscale ( "<STR_LIT>" ) <EOL> plt . ylabel ( "<STR_LIT>" ) <EOL> plt . show ( ) <EOL> plt . savefig ( os . path . join ( results_dir , "<STR_LIT>" ) ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> run = <NUM_LIT:7> <EOL> start_train_year = <NUM_LIT> <EOL> end_train_year = <NUM_LIT> <EOL> N_names = <NUM_LIT:100> <EOL> state = None <EOL> downsample = <NUM_LIT:50> <EOL> train , test , full_train , full_test = load_data ( start_train_year = start_train_year , <EOL> end_train_year = end_train_year , <EOL> N_names = N_names , <EOL> continental = True , DC = False , <EOL> train_state = state , <EOL> downsample = downsample ) <EOL> k_year = Matern52 ( input_dim = <NUM_LIT:1> , active_dims = [ <NUM_LIT:0> ] , variance = <NUM_LIT:10> , lengthscale = <NUM_LIT:30> ) <EOL> k_latlon = RBF ( input_dim = <NUM_LIT:2> , active_dims = [ <NUM_LIT:1> , <NUM_LIT:2> ] , variance = <NUM_LIT:0.1> , lengthscale = <NUM_LIT:5> ) <EOL> kernel = k_year * k_latlon <EOL> from pgmult . internals . utils import mkdir <EOL> results_dir = os . path . join ( "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" % run ) <EOL> mkdir ( results_dir ) <EOL> results_file = os . path . join ( "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" % run , "<STR_LIT>" ) <EOL> static_model = pgmult . distributions . IndependentMultinomialsModel ( train . data [ train . years == end_train_year - <NUM_LIT:1> ] ) <EOL> static_pll = compute_static_pred_ll ( static_model , full_test ) <EOL> print ( "<STR_LIT>" % ( end_train_year - <NUM_LIT:1> , static_pll ) ) <EOL> results_file = os . path . join ( results_dir , "<STR_LIT>" ) <EOL> if os . path . exists ( results_file ) : <EOL> with gzip . open ( results_file ) as f : <EOL> emp_gp_model , emp_pll , emp_ppi = pickle . load ( f ) <EOL> else : <EOL> emp_gp_model = pgmult . gp . EmpiricalStickBreakingGPModel ( train . K , kernel , D = <NUM_LIT:3> , alpha = <NUM_LIT:0.1> ) <EOL> emp_gp_model . add_data ( get_inputs ( train ) , train . data , optimize_hypers = True ) <EOL> emp_pll , emp_ppi = emp_gp_model . predictive_log_likelihood ( get_inputs ( full_test ) , full_test . data ) <EOL> with gzip . open ( results_file , "<STR_LIT:w>" ) as f : <EOL> pickle . dump ( ( emp_gp_model , emp_pll , emp_ppi ) , f , protocol = - <NUM_LIT:1> ) <EOL> print ( "<STR_LIT>" , emp_pll ) <EOL> emp_pi_train = emp_gp_model . predict ( get_inputs ( train ) ) [ <NUM_LIT:0> ] <EOL> results_file = os . path . join ( results_dir , "<STR_LIT>" ) <EOL> N_samples = <NUM_LIT:20> <EOL> burnin = N_samples // <NUM_LIT:2> <EOL> if os . path . exists ( results_file ) : <EOL> with gzip . open ( results_file ) as f : <EOL> sbm_gp_samples , sbm_gp_plls , sbm_gp_ppis , sbm_gp_timestamps = pickle . load ( f ) <EOL> else : <EOL> gp_model = pgmult . gp . MultinomialGP ( train . K , emp_gp_model . kernel , D = <NUM_LIT:3> ) <EOL> gp_model . add_data ( get_inputs ( train ) , train . data ) <EOL> sbm_gp_samples , _ , sbm_gp_plls , sbm_gp_ppis , sbm_gp_timestamps = fit_gp_multinomial_model ( gp_model , full_test , pi_train = emp_pi_train , N_samples = N_samples ) <EOL> with gzip . open ( results_file , "<STR_LIT:w>" ) as f : <EOL> pickle . dump ( ( sbm_gp_samples [ - <NUM_LIT:1> : ] , sbm_gp_plls , sbm_gp_ppis , sbm_gp_timestamps ) , f , protocol = - <NUM_LIT:1> ) <EOL> sbm_avg_pll = logsumexp ( sbm_gp_plls [ burnin : ] ) - np . log ( N_samples - burnin ) <EOL> print ( "<STR_LIT>" , sbm_avg_pll ) <EOL> N_samples = <NUM_LIT:200> <EOL> results_file = os . path . join ( results_dir , "<STR_LIT>" ) <EOL> ln_gp_model = pgmult . gp . LogisticNormalGP ( train . K , emp_gp_model . kernel , D = <NUM_LIT:3> ) <EOL> ln_gp_model . add_data ( get_inputs ( train ) , train . data ) <EOL> if os . path . exists ( results_file ) : <EOL> with gzip . open ( results_file ) as f : <EOL> lnm_gp_samples , lnm_gp_plls , lnm_gp_ppis , lnm_gp_timestamps = pickle . load ( f ) <EOL> else : <EOL> lnm_gp_samples , _ , lnm_gp_plls , lnm_gp_ppis , lnm_gp_timestamps = fit_gp_multinomial_model ( ln_gp_model , full_test , pi_train = emp_pi_train , N_samples = N_samples ) <EOL> with gzip . open ( results_file , "<STR_LIT:w>" ) as f : <EOL> pickle . dump ( ( lnm_gp_samples [ - <NUM_LIT:1> : ] , lnm_gp_plls , lnm_gp_ppis , lnm_gp_timestamps ) , f , protocol = - <NUM_LIT:1> ) <EOL> lnm_avg_pll = logsumexp ( lnm_gp_plls [ burnin : ] ) - np . log ( N_samples - burnin ) <EOL> print ( "<STR_LIT>" , lnm_avg_pll ) <EOL> print ( "<STR_LIT>" ) <EOL> compute_fraction_top_names ( full_test , [ np . vstack ( ( static_model . pi , static_model . pi ) ) ] ) <EOL> print ( "<STR_LIT>" ) <EOL> compute_fraction_top_names ( full_test , [ emp_ppi ] ) <EOL> print ( "<STR_LIT>" ) <EOL> compute_fraction_top_names ( full_test , sbm_gp_ppis [ burnin : ] ) <EOL> print ( "<STR_LIT>" ) <EOL> compute_fraction_top_names ( full_test , lnm_gp_ppis [ burnin : ] ) <EOL> plot_pred_ll_vs_time ( static_pll , emp_pll , <EOL> sbm_gp_timestamps , sbm_gp_plls , sbm_avg_pll , <EOL> lnm_gp_timestamps , lnm_gp_plls , lnm_avg_pll , <EOL> results_dir = results_dir , <EOL> norm = full_test . data . sum ( ) , <EOL> baseline = static_pll ) </s>
