%pylab inline

import math

import astropy.units as u
import astropy.constants as const

import scipy.interpolate
import scipy.fftpack

# yt is not included in anaconda but installs easily with "pip install yt".
# It is only used to make one plot below, so you skip this if you don't need that plot.
import yt

from astropy.cosmology import Planck13 as fiducial

h0 = fiducial.H0/(100*u.km/u.s/u.Mpc)
print 'Fiducial h0 =',h0

kvec,planck13_Pk = np.loadtxt('camb/Planck13_matterpower_1.dat').transpose()

pk_interpolator = scipy.interpolate.InterpolatedUnivariateSpline(np.log(kvec),np.log(planck13_Pk),k=1)

def pk_func(kvec_hMpc):
    pk = np.zeros_like(kvec_hMpc)
    nonzero = (kvec_hMpc>0)
    pk[nonzero] = np.exp(pk_interpolator(np.log(kvec_hMpc[nonzero])))
    # Units of P(k) are (h/Mpc)**2
    return pk

kR = 8.*kvec
sigma8_window = (3/kR**3*(np.sin(kR)-kR*np.cos(kR)))**2*kvec**3/(2*np.pi**2)

fig = plt.figure(figsize=(6,4))
plt.fill_between(kvec,planck13_Pk,edgecolor='blue',lw=1,facecolor=(0.,0.,1.,0.25))
plt.xlim(kvec[0],kvec[-1])
plt.xlabel('Wavenumber k (h/Mpc)')
plt.ylabel('Power P(k) (h/Mpc)^3',color='blue')
plt.xscale('log')
ax2 = plt.gca().twinx()
ax2.fill_between(kvec,sigma8_window,edgecolor='red',lw=0.5,facecolor=(1.,0.,0.,0.25))
ax2.set_xlim(kvec[0],kvec[-1])
ax2.set_yticklabels([])
ax2.set_yticks([]);
ax2.set_ylabel('Sigma8 Window Function',color='red');
plt.savefig('plots/sigma8window.pdf')

k0 = 0.05 # h/Mpc
r0 = 100 # Mpc/h
powerlaw_Pk = 1e6*kvec/(1+(kvec/k0)**3.5)
wiggles_Pk = powerlaw_Pk*(1 + np.sin(kvec*r0)*np.exp(-(kvec*r0)**2/1e4))
bump1_Pk = 4e4*np.exp(-np.log(kvec/k0)**2)
bump2_Pk = 4e4*np.exp(-np.log(kvec/(2*k0))**2)
noscale_Pk = kvec**-2

plt.plot(kvec/k0,planck13_Pk,'k-')
plt.plot(kvec/k0,powerlaw_Pk,'r-')
plt.plot(kvec/k0,wiggles_Pk,'b-')
plt.plot(kvec/k0,bump1_Pk,'g-')
plt.plot(kvec/k0,bump2_Pk,'g--')
plt.plot(kvec/k0,noscale_Pk,'magenta')
plt.xlim(kvec[0]/k0,kvec[-1]/k0)
plt.ylim(1e-3,1e5)
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Wavenumber k/k0')
plt.ylabel('Power P(k) (h/Mpc)^3')
plt.grid();

np.savetxt('alt/powerlaw_Pk.dat',np.vstack((kvec,powerlaw_Pk)).T)
np.savetxt('alt/wiggles_Pk.dat',np.vstack((kvec,wiggles_Pk)).T)
np.savetxt('alt/bump1_Pk.dat',np.vstack((kvec,bump1_Pk)).T)
np.savetxt('alt/bump2_Pk.dat',np.vstack((kvec,bump2_Pk)).T)
np.savetxt('alt/noscale_Pk.dat',np.vstack((kvec,noscale_Pk)).T)

rvec,planck13_xi = np.loadtxt('camb/Planck13_xi.dat').transpose()
rvec,powerlaw_xi = np.loadtxt('alt/powerlaw_xi.dat').transpose()
rvec,wiggles_xi = np.loadtxt('alt/wiggles_xi.dat').transpose()
rvec,bump1_xi = np.loadtxt('alt/bump1_xi.dat').transpose()
rvec,bump2_xi = np.loadtxt('alt/bump2_xi.dat').transpose()
rvec,noscale_xi = np.loadtxt('alt/noscale_xi.dat').transpose()
planck13_xi /= 2*np.pi**2
powerlaw_xi /= 2*np.pi**2
wiggles_xi /= 2*np.pi**2
bump1_xi /= 2*np.pi**2
bump2_xi /= 2*np.pi**2
noscale_xi /= 2*np.pi**2

plt.plot(rvec*k0,rvec**2*planck13_xi,'k-')
plt.plot(rvec*k0,rvec**2*powerlaw_xi,'r-')
plt.plot(rvec*k0,rvec**2*wiggles_xi,'b-')
plt.plot(rvec*k0,rvec**2*bump1_xi,'g-')
plt.plot(rvec*k0,rvec**2*bump2_xi,'g--')
plt.plot(rvec*k0,rvec**2*noscale_xi,'magenta')
plt.xlim(0.,rvec[-1]*k0)
plt.ylim(-150.,500.)
plt.xlabel('Separation r*k0')
plt.ylabel('Correlation function r^2*xi(r) (Mpc/h)^2');

planck13_slice = np.loadtxt('camb/planck13_slice.dat')
powerlaw_slice = np.loadtxt('alt/powerlaw_slice.dat')
wiggles_slice = np.loadtxt('alt/wiggles_slice.dat')
bump1_slice = np.loadtxt('alt/bump1_slice.dat')
bump2_slice = np.loadtxt('alt/bump2_slice.dat')
noscale_slice = np.loadtxt('alt/noscale_slice.dat')

def plot_Pk(data,label):
    plt.plot(kvec/k0,data,'k-')
    visible = (1e-2 < kvec/k0) & (kvec/k0 < 1e2)
    pmax = np.max(data[visible])
    plt.xlim(1e-2,1e2)
    plt.ylim(2e-4*pmax,2*pmax)
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Wavenumber $k/k_{0}$')
    plt.ylabel('Power $P(k)$')
    plt.gca().get_yaxis().set_ticklabels([])
    plt.grid()
    plt.annotate(label,xy=(0.85,0.85),xycoords='axes fraction',fontsize='xx-large')
    
def plot_xi(data,label):
    plt.plot(k0*rvec,rvec**2*data,'k-')
    plt.xlim(0.,rvec[-1]*k0)
    if np.min(data) >= 0:
        plt.ylim(-0.05*np.max(rvec**2*data),None)
    plt.xlabel('Separation $k_{0} r$')
    plt.ylabel('Correlation $r^{2} \\xi(r)$')
    plt.gca().get_yaxis().set_ticklabels([])
    plt.gca().axhline(0., linestyle='--', color='k')
    plt.annotate(label,xy=(0.85,0.85),xycoords='axes fraction',fontsize='xx-large')
    
# Display a 2D density field using a color map where 0 = white, >0 is red, <0 is blue.
# The input data is assumed to have a mean near zero.
def slice_plot(data,clip_percent=1.,**kwargs):
    lo,hi = np.percentile(data,(clip_percent,100.-clip_percent))
    assert lo < 0. and hi > 0.
    cut = 0.5*(hi-lo)
    plt.imshow(data,cmap='RdBu_r',vmin=lo,vmax=hi,**kwargs)
    
def plot_slice(data,label):
    slice_plot(data,extent=(0.,5.,0.,5.))
    plt.xlabel('$k_{0}x$')
    plt.ylabel('$k_{0}y$')
    plt.annotate(label,xy=(0.85,0.85),xycoords='axes fraction',fontsize='xx-large')

def match1(save_name=None):
    fig = plt.figure(figsize=(7.5,7.5))
    # planck13 P(k)
    plt.subplot(3,3,1)
    plot_Pk(planck13_Pk,'A')
    # bump2 P(k)
    plt.subplot(3,3,4)
    plot_Pk(bump2_Pk,'B')
    # bump1 P(k)
    plt.subplot(3,3,7)
    plot_Pk(bump1_Pk,'C')
    # bump2 xi(r)
    plt.subplot(3,3,2)
    plot_xi(bump2_xi,'A')
    # bump1 xi(r)
    plt.subplot(3,3,5)
    plot_xi(bump1_xi,'B')
    # planck13 xi(r)
    plt.subplot(3,3,8)
    plot_xi(planck13_xi,'C')
    # bump1 realization
    plt.subplot(3,3,3)
    plot_slice(bump1_slice,'A')
    # planck13 realization
    plt.subplot(3,3,6)
    plot_slice(planck13_slice,'B')
    # bump2 realization
    plt.subplot(3,3,9)
    plot_slice(bump2_slice,'C')
    #
    plt.tight_layout()
    if save_name:
        plt.savefig(save_name)
match1('plots/match1.pdf')

def match2(save_name=None):
    fig = plt.figure(figsize=(7.5,7.5))
    # noscale P(k)
    plt.subplot(3,3,1)
    plot_Pk(noscale_Pk,'A')
    # powerlaw P(k)
    plt.subplot(3,3,4)
    plot_Pk(powerlaw_Pk,'B')
    # wiggles P(k)
    plt.subplot(3,3,7)
    plot_Pk(wiggles_Pk,'C')
    # powerlaw xi(r)
    plt.subplot(3,3,2)
    plot_xi(powerlaw_xi,'A')
    # noscale xi(r)
    plt.subplot(3,3,5)
    plot_xi(noscale_xi,'B')
    # wiggles xi(r)
    plt.subplot(3,3,8)
    plot_xi(wiggles_xi,'C')
    # noscale realization
    plt.subplot(3,3,3)
    plot_slice(noscale_slice,'A')
    # powerlaw realization
    plt.subplot(3,3,6)
    plot_slice(powerlaw_slice,'B')
    # wiggles realization
    plt.subplot(3,3,9)
    plot_slice(wiggles_slice,'C')
    #
    plt.tight_layout()
    if save_name:
        plt.savefig(save_name)
match2('plots/match2.pdf')

def generate_field(spacing_Mpch=1.,num_grid=64,seed=2015):
    dk = 2*np.pi/(spacing_Mpch*num_grid) # in h/Mpc
    kvec = 2*np.pi*np.fft.fftfreq(num_grid,d=spacing_Mpch) # in h/Mpc
    kx,ky,kz = np.meshgrid(kvec,kvec,kvec,sparse=True)
    knorm = np.sqrt(kx**2 + ky**2 + kz**2)
    del kx,ky,kz # We don't need these large arrays any more
    print 'Sampling P(k) for k = (%.3f - %.3f) h/Mpc' % (np.min(knorm[knorm>0]),np.max(knorm))
    # Calculate the variance in each k-space cell of the real and imaginary parts of the k-space delta field.
    variance = pk_func(knorm)*dk**3/(8*np.pi**3)
    del knorm # We don't need this large array any more
    # Generate a complex-valued random number for each cell, with its real and imaginary parts
    # sampled from independent zero-mean unit Gaussians, scaled to the correct variance.
    # The extra factor of num_grid**3 is due to the normalization of scipy.fftpack inverse transforms.
    generator = np.random.RandomState(seed)
    delta_k = num_grid**3*np.sqrt(variance)*(generator.normal(size=(num_grid,num_grid,2*num_grid)).view(dtype=complex))
    del variance # We don't need this large array any more
    # Transform the k-space complex delta field to an r-space complex delta field.
    delta_r = scipy.fftpack.ifftn(delta_k,overwrite_x=True).view(dtype=float).reshape(num_grid,num_grid,num_grid,2)
    # Since we treat the k-space delta field as a general complex field, its inverse FT is not real.
    # We could save some memory by building in the necessary k-space symmetries and using a special iFFT algorighm,
    # but instead we return a pair of uncorrelated real fields (as views into the complex r-space field)
    return delta_r[:,:,:,0],delta_r[:,:,:,1]

delta_r1,delta_r2 = generate_field(spacing_Mpch=8./15.,num_grid=256);

def create_slice(delta_r,num_zavg):
    return np.sum(delta_r[:,:,:num_zavg],axis=-1)/num_zavg

slice1 = create_slice(delta_r1,20)
slice2 = create_slice(delta_r2,20)

tracer_slice1 = np.loadtxt('camb/tracer_slice1.dat')
tracer_slice2 = np.loadtxt('camb/tracer_slice2.dat')

np.var(slice1),np.var(slice2),np.var(tracer_slice1),np.var(tracer_slice2)

def sample_tracers(slice_data,num_tracers,spacing,bias=1.5,seed=2015):
    print 'Using bias = %.3f' % bias
    nx,ny = slice_data.shape
    print 'mean tracer density = %.5f / (units of spacing)**2' % (float(num_tracers)/(nx*ny*spacing))
    # Set the overall mean number of tracers averaged over all cells.
    overall_mean = float(num_tracers)/slice_data.size
    # Calculate the mean number of tracers in each cell.
    cell_mean = overall_mean*(1 + bias*slice_data)
    # Clip any negative cell means.
    print 'Clipped %d / %d cells with mean < 0.' % (np.count_nonzero(cell_mean < 0),nx*ny)
    cell_mean[cell_mean < 0] = 0.
    # Generate the actual number of tracers in each cell.
    generator = np.random.RandomState(seed)
    cell_count = generator.poisson(cell_mean)
    # Clip any cells with more than 1 tracer.
    clipped = (cell_count > 1)
    print 'Clipped %d / %d cells with > 1 tracer.' % (np.count_nonzero(clipped),nx*ny)
    cell_count[clipped] = 1
    # Create the coordinate grid.
    xgrid = np.linspace(-nx+1,nx-1,nx,endpoint=True)*spacing/2.
    ygrid = np.linspace(-ny+1,ny-1,ny,endpoint=True)*spacing/2.
    x,y = np.meshgrid(xgrid,ygrid)
    filled = (cell_count > 0)
    print 'generated %d tracers' % np.count_nonzero(filled)
    return x[filled],y[filled]

tracer1_x,tracer1_y = sample_tracers(tracer_slice1,1000,8./15./h0)
tracer2_x,tracer2_y = sample_tracers(tracer_slice2,1000,8./15./h0)

data = dict(Density = delta_r1)
size = 128*8./15./h0
bbox = np.array([[-size,+size],[-size,+size],[-size,+size]])
ds = yt.load_uniform_grid(data,delta_r1.shape,length_unit = 'Mpc',bbox = bbox,nprocs = 4)

slc = yt.SlicePlot(ds,'z',['Density'])
slc.set_cmap('Density','RdBu_r')
slc.set_log('Density',False)
slc.show()

rms = np.std(delta_r1)
print 'RMS',rms
tf = yt.ColorTransferFunction((-0.5*rms,+0.5*rms))
#tf.add_layers(N = 3,alpha = [0.1]*5)#,colormap = 'RdBu_r')
thickness = 2*(0.25*rms)**2
tf.add_gaussian(-0.5*rms,thickness,[0.,0.,1.,0.25])
tf.add_gaussian(0.,thickness,[1.,1.,1.,0.10])
tf.add_gaussian(0.5*rms,thickness,[1.,0.,0.,0.25])
center = (ds.domain_right_edge + ds.domain_left_edge)/2.0
look = np.array([1.0, 1.0, 1.0])
width = 20. # Mpc
npix = [1024,768] # 4:3 frame
cam = ds.camera(center,look,width,npix,tf,fields = ['Density'],log_fields = [False],no_ghost=False)
im = cam.snapshot('plots/volumetric.png')

us_dime = 17.95*u.mm
us_quarter = 24.26*u.mm
euro = 23.25*u.mm

def draw_frame(coin = us_quarter, radius_Mpch = 8., z0 = 1.2):
    fig = plt.figure(figsize=(7.5,7.5))
    axes = plt.subplot(1,1,1)
    # Initialize a plot with seperate x,y axes on both sides.
    axes.set_xlabel('Transverse Comoving Distance (Mpc)')
    axes.set_ylabel('Line-of-sight Comoving Distance (Mpc)')
    # Set temporary limits so we can fix the aspect ratio.
    plt.xlim(-15,+15)
    plt.ylim(-15,+15)
    axes.set_aspect('equal',adjustable='box')
    # Get the axes sizes in inches. Remember to trigger updates via draw() before
    # calling get_window_extent().
    plt.draw()
    bbox = axes.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    assert abs(bbox.width - bbox.height) < 1e-6
    print 'Axes physical size = %.3f inches' % bbox.width
    # Calculate the axes length in Mpc so that the coin diameter is exactly twice
    # the specified radius in Mpc/h.
    coin_diameter_inches = coin.to(u.imperial.inch).value
    print 'With h0 = %.3f, %f Mpc/h = %.3f Mpc' % (h0,radius_Mpch,radius_Mpch/h0)
    comoving_size = (2*radius_Mpch/h0)*(bbox.width/coin_diameter_inches)
    print 'Axes comoving size = %.3f Mpc' % comoving_size
    plt.xlim(-0.5*comoving_size,+0.5*comoving_size)
    plt.ylim(-0.5*comoving_size,+0.5*comoving_size)
    # Calculate the corresponding range of angular separations (in degrees) at redshift z0.
    angular_size = math.degrees(comoving_size*u.Mpc/((1+z0)*fiducial.angular_diameter_distance(z=z0)))
    print 'Axes angular size = %.3f degrees' % angular_size
    # Calculate the corresponding redshift range centered at z0.
    redshift_size = comoving_size*u.Mpc/((const.c/fiducial.H(z0)).to(u.Mpc))
    print 'Axes redshift size = %.5f' % redshift_size
    # Add secondary axes (which requires resetting the original axes).
    alt_x_axis = axes.twiny()
    alt_y_axis = axes.twinx()
    axes.set(adjustable='box-forced',xlim=(-0.5*comoving_size,+0.5*comoving_size),
        ylim=(-0.5*comoving_size,+0.5*comoving_size),aspect=1.)
    alt_x_axis.set(adjustable='box-forced',xlim=(-0.5*angular_size,+0.5*angular_size),
        ylim=(-0.5*comoving_size,+0.5*comoving_size),
        aspect=angular_size/comoving_size)
    alt_x_axis.set_xlabel('Angular Separation (degrees)')
    alt_y_axis.set(adjustable='box-forced',xlim=(-0.5*comoving_size,+0.5*comoving_size),
        ylim=(z0-0.5*redshift_size,z0+0.5*redshift_size),
        aspect=comoving_size/redshift_size)
    alt_y_axis.set_ylabel('Redshift')
    # Calculate cosmic-variance limit.
    nmax = comoving_size**2/(np.pi*(radius_Mpch/h0)**2)
    print 'Maximum number of independent samples is %.1f' % nmax
    return axes,comoving_size

def make_handout(x,y,style,save_name=None,coin=us_quarter):
    ax,size = draw_frame(radius_Mpch = 8.,coin = coin)
    ax.plot(x,y,style)
    ax.add_artist(plt.Circle((0,0),8./h0,edgecolor='lightgray',facecolor='none'))
    # Count number of visible galaxies.
    half_size_Mpch = 0.5*size*h0
    visible = (x > -half_size_Mpch) & (x < +half_size_Mpch) & (y > -half_size_Mpch) & (y < +half_size_Mpch)
    print 'Number of visible tracers %d' % np.count_nonzero(visible)
    if save_name:
        plt.savefig(save_name)

make_handout(tracer1_x,tracer1_y,'r+','plots/slice1.pdf')

make_handout(tracer2_x,tracer2_y,'bx','plots/slice2.pdf')

make_handout(tracer1_x,tracer1_y,'r+','plots/slice1_euro.pdf',coin = euro)

make_handout(tracer2_x,tracer2_y,'bx','plots/slice2_euro.pdf',coin = euro)

figure(figsize=(8,4))
plt.subplot(1,2,1)
slice_plot(noscale_slice,clip_percent=1.)
plt.axis('off')
plt.subplot(1,2,2)
slice_plot(planck13_slice,clip_percent=1.)
plt.axis('off')
plt.tight_layout()
plt.savefig('plots/structure.png')

zNL,kNLlo,kNLhi = np.loadtxt('camb/nonlinearity.dat').transpose()

plt.fill_between(zNL,2*np.pi/kNLlo/h0,125.,facecolor='green',alpha=0.15,lw=0)
plt.fill_between(zNL,2*np.pi/kNLhi/h0,2*np.pi/kNLlo/h0,facecolor='red',alpha=0.25,edgecolor='k',lw=2,linestyle='dashed')
plt.fill_between(zNL,2*np.pi/kNLhi/h0,0.,facecolor='red',alpha=0.5,edgecolor='k',lw=2)
plt.ylim(0.,125.)
plt.xlabel('Redshift z')
plt.ylabel('Scale (Mpc/h)')
plt.grid()
plt.tight_layout()
plt.savefig('plots/nonlinearity.pdf')

mean_tracer_density = 0.01939 # per Mpc^2
avg_under_coin = np.pi*(8./h0)**2*mean_tracer_density
print 'Average galaxies under a coin mu = %.3f' % avg_under_coin
generator = np.random.RandomState(123)
samples = generator.poisson(avg_under_coin,100)
bins = np.arange(26)-0.5
plt.hist(samples,bins=bins,histtype='stepfilled',color='gray',alpha=0.5)
plt.xlabel('Galaxies / Quarter')
plt.xlim(bins[0],bins[-1])
plt.grid()
plt.tight_layout()
print 'mean = %.3f, std.dev. = %.3f' % (np.mean(samples),np.std(samples))
plt.savefig('plots/histogram.pdf')

mu = np.mean(samples)
delta_coin = samples/mu-1.
bins = (np.arange(26)-0.5)/mu-1.
plt.hist(delta_coin,bins=bins,histtype='stepfilled',color='gray',alpha=0.5)
plt.xlabel('Galaxy number density contrast $\delta_{g}(r)$')
plt.xlim(bins[0],bins[-1])
plt.grid()
plt.tight_layout()
print 'mean = %.3f, std.dev. = %.3f' % (np.mean(delta_coin),np.std(delta_coin))
plt.savefig('plots/delta.pdf')

correlation_coef = -0.5
cov_matrix = np.array([[1.,correlation_coef],[correlation_coef,1.]])*np.var(samples)
mu_vector = np.array([1.,1.])*avg_under_coin
generator = np.random.RandomState(123)
x,y = generator.multivariate_normal(mu_vector,cov_matrix,1000).transpose()
fig = plt.figure(figsize=(9,8))
plt.subplot(2,2,1)
plt.scatter(x,y,color='b',lw=0)
plt.gca().set_aspect(1.)
plt.xlim(0.,20.)
plt.ylim(0.,20.)
plt.xlabel('Galaxies / quarter')
plt.ylabel('Galaxies / quarter')
plt.subplot(2,2,2)
plt.hist(x,bins=25,histtype='step',color='b')
plt.hist(y,bins=25,histtype='step',color='b',linestyle='dashed')
plt.xlabel('Galaxies / quarter')
plt.subplot(2,2,3)
x = x/mu-1.
y = y/mu-1.
plt.scatter(x,y,color='b',lw=0)
plt.gca().set_aspect(1.)
plt.xlim(-1.,+1.)
plt.ylim(-1.,+1.)
plt.grid()
plt.xlabel('$\delta_g(r_1)$')
plt.ylabel('$\delta_g(r_2)$')
plt.subplot(2,2,4)
plt.hist(x,bins=25,histtype='step',color='b')
plt.hist(y,bins=25,histtype='step',color='b',linestyle='dashed')
plt.xlabel('$\delta_g(r_1)$')
plt.xlim(-1.,+1.)
plt.grid()
plt.tight_layout();
plt.savefig('plots/correlation.pdf')

plt.scatter(x,y,color='b',lw=0)
plt.gca().set_aspect(1.)
plt.xlim(-1.,+1.)
plt.ylim(-1.,+1.)
plt.gca().set_aspect(1.)
plt.grid()
plt.xlabel('$\delta_g(r_1)$')
plt.ylabel('$\delta_g(r_2)$')
plt.tight_layout();
plt.savefig('plots/scatter.pdf')

plt.hist(x,bins=25,histtype='stepfilled',color='b',alpha=0.5)
plt.xlim(-1.,+1.)
plt.grid()
plt.gca().set_yticklabels([])
plt.tight_layout();
plt.savefig('plots/xhist.pdf')

plt.hist(y,bins=25,histtype='stepfilled',color='b',alpha=0.5)
plt.xlim(-1.,+1.)
plt.grid()
plt.gca().set_yticklabels([])
plt.tight_layout();
plt.savefig('plots/yhist.pdf')

fig = plt.figure(figsize=(9,7))
grid = np.arange(0.,200.,10.) # in Mpc/h
xi_func = scipy.interpolate.InterpolatedUnivariateSpline(rvec,planck13_xi,k=1)
x,y = np.meshgrid(grid,grid)
dr = np.fabs(x-y)
# Use (1+dr)**2 instead of dr**2 weighting so value is not zero along diagonal.
Cr = (1.+dr)**2*xi_func(dr.flat).reshape(20,20)
plt.matshow(Cr,cmap='RdBu_r',extent=(0.,grid[-1],0.,grid[-1]),vmin=-10.,vmax=+10.)
plt.xlabel('$x$ (Mpc/h)')
plt.ylabel('$y$ (Mpc/h)')
plt.colorbar(label='$\\xi(|r_2 - r_1|)$',pad=0.01)
plt.savefig('plots/Cr.pdf');

# The built-in 'Reds' colormap doesn't go all the way to white, so we make our own here.
from matplotlib.colors import LinearSegmentedColormap
cmap_wtr = LinearSegmentedColormap.from_list('wtr',('white','red'));
klo,khi = 0.01,1.
k = np.logspace(np.log(klo),np.log(khi),20)
Pk = pk_func(k)
Ck = np.diag(np.log(Pk))
fig = plt.figure(figsize=(9,7))
# This command generates a UserWarning the first time it is run, which seems to be harmless.
plt.matshow(Ck,cmap=cmap_wtr,extent=(klo,khi,klo,khi),vmin=0.,vmax=10.)
plt.xscale('log')
plt.yscale('log')
plt.xlabel('$k_x$ (h/Mpc)')
plt.ylabel('$k_y$ (h/Mpc)')
plt.colorbar(label='$log P(k)$',pad=0.01)
plt.savefig('plots/Ck.pdf');

grid = np.linspace(-4.,4.,101,endpoint=True)
fig = plt.figure(figsize=(8,4))
plt.subplot(1,2,1)
x0,y0 = 0.5,-1.5
kx,ky = np.meshgrid(grid,grid)
kmode = np.cos(kx*x0 + ky*y0)
plt.scatter(x0,y0)
plt.xlim(grid[0],grid[-1])
plt.ylim(grid[0],grid[-1])
plt.xlabel('$x$')
plt.ylabel('$y$')
plt.grid()
plt.subplot(1,2,2)
plt.imshow(kmode,cmap='RdBu_r',vmin=-1.,vmax=+1.,extent=(grid[0],grid[-1],grid[0],grid[-1]),origin='lower')
plt.xlabel('$k_x$')
plt.ylabel('$k_y$')
plt.grid()
plt.tight_layout();
plt.savefig('plots/modes.pdf')