def plot_spatial_distribution(train, samples, name="Ethan", year=2000):
    # Extract data from samples
    # Extract samp[les
    mus = np.array([s[0] for s in samples])
    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
    psis += mus[0][None,None,:]
    pis = np.array([psi_to_pi(psi_sample) for psi_sample in psis])

    # Extract single name, year data
    data = pis[-1, train.years==year, train.names == name.lower()]
    lons = train.lon[train.years == year]
    lats = train.lat[train.years == year]

    fig = plt.figure(figsize=(3,3))
    ax = fig.add_subplot(111, aspect="equal")

    from mpl_toolkits.basemap import Basemap
    m = Basemap(width=6000000, height=3500000,
                resolution='l',projection='stere',
                lat_ts=50,lat_0=40,lon_0=-100.,
                ax=ax)
    land_color  = [.98, .98, .98]
    water_color = [.75, .75, .75]
    # water_color = [1., 1., 1.]
    m.fillcontinents(color=land_color, lake_color=water_color)
    m.drawcoastlines()
    m.drawstates()
    m.drawcountries()
    m.drawmapboundary(fill_color=water_color)

    # Convert data lons and data lats to map coordinates
    dx, dy = m(lons, lats)

    # Interpolate at a grid of points
    glons, glats = m.makegrid(100, 100)
    gx, gy = m(glons, glats)
    M = gx.size

    # Interpolate
    from scipy.interpolate import griddata
    gdata = griddata(np.hstack((dx[:,None], dy[:,None])),
                     data,
                     np.hstack((gx.reshape((M,1)), gy.reshape((M,1)))),
                     method="cubic")
    gdata = gdata.reshape(gx.shape)

    # Plot the contour
    cs = ax.contour(gx, gy, gdata, 15, cmap="Reds", linewidth=2)
    plt.title("%s (%d)" % (name, year))

    from mpl_toolkits.axes_grid1 import make_axes_locatable
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="3%", pad=0.05)
    cbar = plt.colorbar(cs, cax=cax)
    cbar.set_label("Probability", labelpad=10)

    plt.subplots_adjust(left=0.05, bottom=0.1, top=0.9, right=0.85)
    fig.savefig("%s_%d_geo.pdf" % (name.lower(), year))

    return fig, ax, m
dataset/ETHPy150Open HIPS/pgmult/experiments/census_gp.py/plot_spatial_distribution/original