<jupyter_start><jupyter_text>Experiment 1 - Univariate Gaussian Model  HMC Config<jupyter_code># sampling params
NUM_CHAINS = 5
NUM_SAMPLES = 1000
NUM_BURNIN_STEPS = 10000
# define step size adaptation
stepsize_adaptation_kernel = DualAveragingStepSizeAdaptation(num_adaptation_steps=int(NUM_BURNIN_STEPS*0.8))

kernel = NoUTurnSampler(
    step_size=0.01, 
    max_tree_depth=5,
    stepsize_adaptation_kernel=stepsize_adaptation_kernel
)<jupyter_output><empty_output><jupyter_text>VI Config<jupyter_code>NUM_STEPS = 10000
SAMPLE_SIZE = 50
optimizer = tf.optimizers.Adam
init_lr = {
    'meanfield advi': 5e-3, 
    'affine flow': 5e-3,
    'maf': 5e-3,
    'cnf': 1e-4,
    'hnf(1)': 1e-3,
    'hnf(2)': 1e-3,
    'hnf(3)': 1e-3,
    'hnf(5)': 1e-3,
}
func_dict = {
    "mean": np.mean,
    "sd": np.std,
    "min": lambda x: np.percentile(x, 0),
    "hdi 3%": lambda x: np.percentile(x, 3),
    "mode": lambda x: np.percentile(x, 50),
    "hdi 97%": lambda x: np.percentile(x, 97),
    "max": lambda x: np.percentile(x, 100)
}
def get_continuous_flow_bijector(unconstrained_event_dims):
    state_fn = tfk.Sequential()
    state_fn.add(tfk.layers.Dense(128, activation=tfk.activations.tanh))
    state_fn.add(tfk.layers.Dense(128, activation=tfk.activations.tanh))
    state_fn.add(tfk.layers.Dense(unconstrained_event_dims))
    state_fn.build((None, unconstrained_event_dims+1))
    state_time_derivative_fn = lambda t, state: state_fn(tf.concat([tf.fill((state.shape[0],1), t), state], axis=-1))
    return tfb.FFJORD(state_time_derivative_fn, 
                      ode_solve_fn=tfp.math.ode.DormandPrince(first_step_size=0.1).solve, 
                      trace_augmentation_fn=tfb.ffjord.trace_jacobian_hutchinson)
    
def get_hamiltonian_flow_bijector(unconstrained_event_dims, num_flows):
    return tfb.Chain([
        HamiltonianFlow(
            event_dims=unconstrained_event_dims,
            symplectic_integrator=LeapfrogIntegrator(), 
            step_sizes=tf.Variable(0.1), 
            num_integration_steps=2,
            hidden_layers=[128, 128]
        ) for _ in range(num_flows)
    ])

def get_masked_autoregressive_flow_bijector(unconstrained_event_dims):
    return tfb.MaskedAutoregressiveFlow(
        shift_and_log_scale_fn=tfb.AutoregressiveNetwork(params=2, hidden_units=[128, 128], activation='tanh')
    )

def get_affine_flow_bijector(unconstrained_event_dims):
    return AffineFlow(unconstrained_event_dims)


def get_posterior_lift_distribution(unconstrained_event_dims):
    scale_fn = make_scale_fn(unconstrained_event_dims, hidden_layers=[128,128])
    shift_fn = make_shift_fn(unconstrained_event_dims, hidden_layers=[128,128])
    return lambda q: tfd.MultivariateNormalDiag(loc=shift_fn(q), scale_diag=scale_fn(q))<jupyter_output><empty_output><jupyter_text>Define Model<jupyter_code>priors = OrderedDict(
    loc = tfd.Normal(loc=0., scale=10.),
    scale = tfd.HalfNormal(scale=10.)
)

likelihood = lambda loc, scale: tfd.Normal(
    loc=loc, 
    scale=scale
)
model = Model(priors=priors, likelihood=likelihood)<jupyter_output><empty_output><jupyter_text>Generate Datasets<jupyter_code>num_datasets = 10
num_datapoints = 100
true_params = {}
datasets = {}

for i in range(1, num_datasets+1):
    dist = model.get_joint_distribution(num_samples=num_datapoints)
    loc, scale, y = dist.sample().values()
    
    true_params['dataset {}'.format(i)] = to_ordered_dict(
        model.param_names, [loc.numpy(), scale.numpy()]
    ) 
    datasets['dataset {}'.format(i)] = make_dataset_from_df(
        pd.DataFrame({'y': y}), target_names=['y'], format_features_as='dict'
    )<jupyter_output><empty_output><jupyter_text>Run Inferences and Collect Results<jupyter_code>results = {}
for i, (name, dataset) in enumerate(datasets.items()):
    print('Starting with {}\n'.format(name))
    features, targets = list(dataset.batch(dataset.cardinality()).take(1))[0]
    
    results[name]={
        'data': {
            'targets': targets.numpy(),
            'true params': true_params[name],
        }}

    ################################################################################################################################
    
    # MCMC
    print('Start MCMC...')

    mcmc = MCMC(model=model, dataset=dataset, transition_kernel=kernel)
    
    print('Run NUTS:')
    start = time.time()
    mcmc_result = mcmc.fit(
        num_chains=NUM_CHAINS, 
        num_samples=NUM_SAMPLES, 
        num_burnin_steps=NUM_BURNIN_STEPS,
        progress_bar=True,
    )
    run_time = round(time.time() - start, 2)
    print('Finished in {}s\n'.format(run_time))

    post_pred_dist = model.get_posterior_predictive_distribution(
        posterior_samples=mcmc_result.samples,
    )
    posterior_predictive_samples = post_pred_dist.sample(1000)['y'].numpy()

    posterior_samples = to_ordered_dict(model.param_names, tf.nest.map_structure(lambda v: np.swapaxes(v.numpy(), 0, 1), mcmc_result.samples))
    map_trace = lambda v: np.swapaxes(v.numpy(), 0, 1) if tf.rank(v) >= 2 else v.numpy()
    traces = {k: map_trace(v) if not tf.nest.is_nested(v) else tf.nest.map_structure(map_trace, v) for k,v in mcmc_result.trace.items()}    
    summary = az.summary(posterior_samples, round_to=2, kind='diagnostics', stat_funcs=func_dict, extend=True)
    cols = summary.columns.tolist()
    summary = summary[cols[-len(func_dict):] + cols[:-len(func_dict)]]
    
    mcmc_results = {
        'posterior samples': posterior_samples,
        'traces': traces,
        'posterior predictive samples': posterior_predictive_samples,
        'acceptance ratios per chain': mcmc_result.accept_ratios.numpy(),
        'summary': summary,
        'run-time': run_time
    }
    
    results[name]['mcmc'] = mcmc_results
    
    print('Finished MCMC! \n')
    ################################################################################################################################
    
    # VI
    unconstrained_event_dims = model.flat_unconstrained_param_event_ndims    

    surrogate_posteriors = {
        'meanfield advi': ADVI(model, mean_field=True), 
        'affine flow': NormalizingFlow(model, flow_bijector=get_affine_flow_bijector(unconstrained_event_dims)),
        'maf': NormalizingFlow(model, flow_bijector=get_masked_autoregressive_flow_bijector(unconstrained_event_dims)),
        'cnf': NormalizingFlow(model, flow_bijector=get_continuous_flow_bijector(unconstrained_event_dims)),
        'hnf(1)': NormalizingFlow(
            model, flow_bijector=get_hamiltonian_flow_bijector(unconstrained_event_dims, num_flows=1), 
            extra_ndims=unconstrained_event_dims, posterior_lift_distribution=get_posterior_lift_distribution(unconstrained_event_dims)),
        'hnf(2)': NormalizingFlow(
            model, flow_bijector=get_hamiltonian_flow_bijector(unconstrained_event_dims, num_flows=2), 
            extra_ndims=unconstrained_event_dims, posterior_lift_distribution=get_posterior_lift_distribution(unconstrained_event_dims)),
        'hnf(3)': NormalizingFlow(
            model, flow_bijector=get_hamiltonian_flow_bijector(unconstrained_event_dims, num_flows=3), 
            extra_ndims=unconstrained_event_dims, posterior_lift_distribution=get_posterior_lift_distribution(unconstrained_event_dims)),
        'hnf(5)': NormalizingFlow(
            model, flow_bijector=get_hamiltonian_flow_bijector(unconstrained_event_dims, num_flows=5), 
            extra_ndims=unconstrained_event_dims, posterior_lift_distribution=get_posterior_lift_distribution(unconstrained_event_dims)),
    }
    
    vi_results = {}
    
    print('Start VI...')
    for surrogate_name, surrogate_posterior in surrogate_posteriors.items():
    
        vi = VI(model, dataset, surrogate_posterior)
        
        print('Run {}:'.format(surrogate_name))
        start = time.time()
        approx_posterior, losses = vi.fit(optimizer=optimizer(init_lr[surrogate_name]), num_steps=NUM_STEPS, sample_size=SAMPLE_SIZE, progress_bar=True)
        run_time = round(time.time() - start, 2)
        print('Finished in {}s\n'.format(run_time))

        print('computing error metrics...')
        start = time.time()

        post_pred_dist = model.get_posterior_predictive_distribution(posterior_distribution=approx_posterior)
        posterior_predictive_samples = post_pred_dist.sample(1000)['y'].numpy()
        
        posterior_samples = tf.nest.map_structure(lambda v: np.expand_dims(v.numpy(),0), approx_posterior.sample(5000))
        summary = az.summary(posterior_samples, stat_funcs=func_dict, extend=False)
        
        # compute Wasserstein distance
        flat_mcmc_samples = tf.nest.map_structure(lambda v: np.expand_dims(v.flatten(), 1), mcmc_results['posterior samples'])
        flat_vi_samples = tf.nest.map_structure(lambda v: np.expand_dims(v.flatten(), 1), posterior_samples)
        distances = tf.nest.map_structure(ot.dist, flat_mcmc_samples, flat_vi_samples)
        W2 = tf.nest.map_structure(
            lambda x,y, dist: np.sqrt(ot.emd2(ot.unif(x.size), ot.unif(y.size), dist)), 
            flat_mcmc_samples, flat_vi_samples, distances
        )

        # compute absolute posterior mean difference to mcmc result
        posterior_mean_error = tf.nest.map_structure(lambda x,y: np.abs(x.mean() - y.mean()), flat_vi_samples, flat_mcmc_samples)
        posterior_sd_error = tf.nest.map_structure(lambda x,y: np.abs(x.std() - y.std()), flat_vi_samples, flat_mcmc_samples)
        comp_time = round(time.time() - start, 2)
        print('Finished computing error metrics in {}s\n\n'.format(comp_time))
        
        vi_result = {
            'losses': losses,
            'final loss': losses[-100:].mean(),
            'learning rate': init_lr[surrogate_name],
            'posterior samples': posterior_samples,
            'posterior predictive samples': posterior_predictive_samples,
            'summary': summary,
            'W2': W2,
            'mean error': posterior_mean_error,
            'sd error': posterior_sd_error,
            'run-time': run_time,
        }
        vi_results[surrogate_name] = vi_result
    print('Finished VI! \n\n\n\n')
    results[name]['vi'] = vi_results
    
with open('./univariate_gaussian_experiment/univariate_gaussian_exp_results.pickle', 'wb') as handle:
    pickle.dump(results, handle, protocol=pickle.HIGHEST_PROTOCOL)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13983050847457626, 'fraction_numerical': 0.01434899845916795, 'mean_word_length': 4.063383715260849, 'pattern_counts': {'":': 7, '<': 21, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '88994', 'license': 'MIT', 'n_tokens_mistral': 3753, 'n_tokens_neox': 3501, 'n_words': 513, 'path': 'univariate_gaussian_experiment/univariate-gaussian-example.ipynb', 'repo_name': 'MaxGrtz/bayesian-inference', 'size': '10384'}