diff --git a/experiments/naturalistic_data.py b/experiments/naturalistic_data.py index c709549..97dd9f2 100644 --- a/experiments/naturalistic_data.py +++ b/experiments/naturalistic_data.py @@ -1,6 +1,3 @@ -from functools import partial - -import jax import click from scipy import stats import jax.numpy as jnp @@ -14,7 +11,7 @@ tfd = tfp.distributions import bayes_ca.inference as core -from bayes_ca.prox_grad import pgd, pgd_jaxopt +from bayes_ca.prox_grad import pgd_jaxopt from bayes_ca.data import naturalistic_data from bayes_ca._utils import _safe_handling_params @@ -133,6 +130,14 @@ def main(seed, data_dir, mu_pri, sigmasq_pri, sigmasq_subj, hazard_prob, max_dur # create train split and initialize globals global_means = jnp.mean(pca_train, axis=0) + _, _, _, smooth_means = core.gaussian_cp_smoother( + global_means, hazard_rates, mu_pri, sigmasq_pri, sigmasq_obs + ) + plt.plot(global_means, label=f"1D input data after FactorAnalysis") + plt.plot(smooth_means, label="smoothed means") + plt.legend() + plt.show() + lps = [] for _ in progress_bar(range(2)): this_key, key = jr.split(key)