def plot_census_results(train, samples, test, test_pis):
    # Extract samp[les
    train_mus = np.array([s[0] for s in samples])
    train_psis = np.array([s[1][0][0] for s in samples])
    # omegas = np.array([s[1][0][1] for s in samples])

    # Adjust psis by the mean and compute the inferred pis
    train_psis += train_mus[0][None,None,:]
    train_pis = np.array([psi_to_pi(psi_sample) for psi_sample in train_psis])
    train_pi_mean = np.mean(train_pis, axis=0)
    train_pi_std = np.std(train_pis, axis=0)

    # Compute test pi mean and std
    test_pi_mean = np.mean(test_pis, axis=0)
    test_pi_std = np.std(test_pis, axis=0)

    # Compute empirical probabilities
    train_pi_emp = train.data / train.data.sum(axis=1)[:,None]
    test_pi_emp = test.data / test.data.sum(axis=1)[:,None]


    # Plot the temporal trajectories for a few names
    names = ["Scott", "Matthew", "Ethan"]
    states = ["NY", "TX", "WA"]
    linestyles = ["-", "--", ":"]

    fig = create_figure(figsize=(3., 3))
    ax1 = create_axis_at_location(fig, 0.6, 0.5, 2.25, 1.75)
    for name, color in zip(names, colors):
        for state, linestyle in zip(states, linestyles):
            train_state_inds = (train.states == state)
            train_name_ind = np.array(train.names) == name.lower()
            train_years = train.years[train.states == state]
            train_mean_name = train_pi_mean[train_state_inds, train_name_ind]
            train_std_name = train_pi_std[train_state_inds, train_name_ind]

            test_state_inds = (test.states == state)
            test_name_ind = np.array(test.names) == name.lower()
            test_years = test.years[test.states == state]
            test_mean_name = test_pi_mean[test_state_inds, test_name_ind]
            test_std_name = test_pi_std[test_state_inds, test_name_ind]

            years = np.concatenate((train_years, test_years))
            mean_name = np.concatenate((train_mean_name, test_mean_name))
            std_name = np.concatenate((train_std_name, test_std_name))

            # Sausage plot
            sausage_plot(years, mean_name, std_name,
                         color=color, alpha=0.5)

            # Plot inferred mean
            plt.plot(years, mean_name,
                     color=color, label="%s, %s" % (name, state),
                     ls=linestyle, lw=2)

            # Plot empirical probabilities
            plt.plot(train.years[train_state_inds],
                     train_pi_emp[train_state_inds, train_name_ind],
                     color=color,
                     ls="", marker="x", markersize=4)

            plt.plot(test.years[test_state_inds],
                     test_pi_emp[test_state_inds, test_name_ind],
                     color=color,
                     ls="", marker="x", markersize=4)

    # Plot a vertical line to divide train and test
    ylim = plt.gca().get_ylim()
    plt.plot((test.years.min()-0.5) * np.ones(2), ylim, ':k', lw=0.5)
    plt.ylim(ylim)

    # plt.legend(loc="outside right")
    plt.legend(bbox_to_anchor=(0., 1.05, 1., .105), loc=3,
               ncol=len(names), mode="expand", borderaxespad=0.,
               fontsize="x-small")

    plt.xlabel("Year")
    plt.xlim(train.years.min(), test.years.max()+0.1)
    plt.ylabel("Probability")

    # plt.tight_layout()
    fig.savefig("census_gp_rates.pdf")

    plt.show()
    plt.pause(0.1)
dataset/ETHPy150Open HIPS/pgmult/experiments/census_gp.py/plot_census_results/original