Predicting on Test Data with Minibatch

Questions v3
1

Hello I am building a hierarchical logistic regression model and using ADVI with mini-batch for training which works great.

However when I go to replace the variables to predict on the test set it does not seem to work and only gives me predictions for a single batch.

The dimensions on the pcc_test at the very bottom are (10000,5000) which is consistent for a single batch. No matter how I configure it I am unable to get it to predict over the entire test set.

Any suggestions are much appreciated!


X_train, y_train = ml_train[target_features].values, ml_train['OutageBool'].values.astype(int)
sub_group_train = ml_train['sub_code'].values.astype(int)
circ_group_train = ml_train['circuit_code'].values.astype(int)

X_test, y_test = ml_test[target_features].values, ml_test['OutageBool'].values.astype(int)
sub_group_test = ml_test['sub_code'].values.astype(int)
circ_group_test = ml_test['circuit_code'].values.astype(int)

batch_size = 5000

coords = {"subs": sub_unique, 
          "circuits": circuit_unique, 
          "obs_id": np.arange(len(y_train)),
        'num_features' : np.array(target_features)}


sub_idx = pm.Minibatch(sub_group_train, batch_size=batch_size)
circ_idx =  pm.Minibatch(circ_group_train, batch_size=batch_size)    

x = pm.Minibatch(X_train, batch_size=batch_size)
y = pm.Minibatch(y_train, batch_size=batch_size)

with pm.Model(coords=coords) as pooled_model:
    
    # Define global parameters that govern the substation coeeficiencts
    global_a = pm.Normal("global_a", mu = -1, sigma=3)
    global_a_unc = pm.Exponential("global_a_unc", 5)

    global_b = pm.Normal("global_b", mu = 1, sigma=3)
    global_b_unc = pm.Exponential("global_b_unc", 5)   
    
    #define substation paramaters
    a_sub = pm.Normal("a_sub", mu = global_a, dims = ( "subs", 'circuits',)) # sigma=global_a_unc,
    a_sub_unc = pm.Exponential("a_sub_unc", global_a_unc, dims= ("subs", 'circuits'))
    
    b_sub = pm.Normal('b_sub', mu = global_b, dims = ("subs",  'circuits', 'num_features')) #sigma=global_b_unc, 
    b_sub_unc = pm.Exponential("b_sub_unc", global_b_unc, dims= ("subs", 'circuits', 'num_features'))    
    
    
    #define circuit paramaters
    a_circ = pm.Normal("a_circ", mu = a_sub, sigma = a_sub_unc, #a_sub_unc[(sub_idx, circ_idx)],
                       dims = ("subs", 'circuits') ) 
    
    b_circ = pm.Normal('b_circ', mu = b_sub, sigma = b_sub_unc, dims = ( "subs", 'circuits', 'num_features') )    
    
    mu = a_circ[sub_idx, circ_idx] + T.sum(b_circ[sub_idx, circ_idx] * x, 1)
    
    p = pm.invlogit(mu)
    
    y = pm.Bernoulli(name='logit', p=p, observed=y, total_size=len(all_circuit_df))   
    

with pooled_model:
    approx = pm.fit(100000, method='advi', )
    advi_trace = approx.sample(10000)
    pooled_idata = az.from_pymc3(advi_trace)

with pooled_model:
    x.set_value(X_test)
    circ_idx.set_value(circ_group_test)
    sub_idx.set_value(sub_group_test)
    
    ppc_test = pm.sample_posterior_predictive(advi_trace, keep_size=False)