%pylab inline

import randomfield
print randomfield.__version__

from randomfield.lensing import *
from randomfield.cosmotools import calculate_power
from randomfield.cosmotools import get_growth_function

"""
Explanation: Lensing Theory Calculations
Examples of using randomfield.lensing functions to perform theoretical calculations related to weak lensing.
End of explanation
"""


from astropy.cosmology import LambdaCDM
flat_model = LambdaCDM(Ob0=0.045, Om0=0.222+0.045, Ode0=0.733, H0=71.0)
open_model = LambdaCDM(Ob0=0.038, Om0=0.186+0.038, Ode0=0.766, H0=77.6)
closed_model = LambdaCDM(Ob0=0.052, Om0=0.256+0.052, Ode0=0.702, H0=66.1)
assert np.allclose((flat_model.Ok0, open_model.Ok0, closed_model.Ok0), (0.0, +0.01, -0.01), atol=1e-4)

"""
Explanation: Define some cosmologies to use for testing.  We use the fiducial models from Table 1 of Weinberg 2012 with $\Omega_k = 0$ (flat) and $\pm 0.01$ (open/closed), which all have essentially the same CMB power spectra.
End of explanation
"""


z = np.linspace(0.0, 2.5, 251)

"""
Explanation: Define the redshift grid to use.  The same grid is used for lensing masses as sources, so you should usually specify a fine grid even if you only need results for a few source redshifts.  Redshifts do not need to be equally spaced.
End of explanation
"""


flat_weights = calculate_lensing_weights(flat_model, z, scaled_by_h=True)
open_weights = calculate_lensing_weights(open_model, z, scaled_by_h=True)
closed_weights = calculate_lensing_weights(closed_model, z, scaled_by_h=True)

"""
Explanation: Calculate the lensing weight function $\omega_E(z, z_{src})$ for each cosmology:
End of explanation
"""


import astropy.units as u
flat_DC = flat_model.comoving_distance(z).to(u.Mpc).value * flat_model.h
open_DC = open_model.comoving_distance(z).to(u.Mpc).value * flat_model.h
closed_DC = closed_model.comoving_distance(z).to(u.Mpc).value * flat_model.h
flat_DA = flat_model.comoving_transverse_distance(z).to(u.Mpc).value * flat_model.h
open_DA = open_model.comoving_transverse_distance(z).to(u.Mpc).value * open_model.h
closed_DA = closed_model.comoving_transverse_distance(z).to(u.Mpc).value * closed_model.h

"""
Explanation: Tabulate the distance functions $D(z)$ and $D_A(z)$ for each cosmology in Mpc/h:
End of explanation
"""


def plot_weights():
    colors = ('r', 'g', 'b', 'y', 'k')
    for j,iz in enumerate(range(50, 251, 50)):
        plt.plot(z, flat_weights[iz]**2 * flat_DA**3, color=colors[j], label='$z_{src}$ = %.2f' % z[iz])
        plt.plot(z, open_weights[iz]**2 * open_DA**3, color=colors[j], ls='--')
        plt.plot(z, closed_weights[iz]**2 * closed_DA**3, color=colors[j], ls=':')
    plt.legend(loc='upper right')
    plt.yscale('log')
    plt.xlabel('Lens redshift $z_{lens}$')
    plt.ylabel('Lensing weight function (h/Mpc)')
    plt.xlim(0,z[-1])
    y_min, y_max = plt.gca().get_ylim()
    plt.ylim(2e-4 * y_max, y_max)
    plt.grid()
    plt.show()

plot_weights()

"""
Explanation: Plot the shear-shear weight function $W_{EE}(z_{lens}, z_{src}) = \omega_E(z_{lens}, z_{src})^2 D_A(z_{lens})^3$ for several source redshifts.  Note that:

Weights are broadly peaked with a maximum value at $z_{lens} \simeq z_{src}/2$.
Weights are larger for larger source redshifts since more distance source experience more lensing.
Weights increase (decrease) slightly in a closed (open) universe.

The different line styles show weights for a flat (solid), open (dashed), or closed (dotted) universe.
End of explanation
"""


ell = np.logspace(1., 3., 101)

"""
Explanation: Define the grid of 2D wavenumbers $\ell$ to use.  Use log-spaced values covering the mostly linear regime. We do not need very fine spacing here since the shear power varies slowly with $\ell$.
End of explanation
"""


plt.plot(z, open_DC, 'r-', label='open')
plt.plot(z, open_DA, 'r--')
plt.plot(z, flat_DC, 'g-', label='flat')
plt.plot(z, closed_DC, 'b-',label='closed')
plt.plot(z, closed_DA, 'b--')
plt.grid()
plt.xlabel('Redshift $z$')
plt.ylabel('Comoving transverse distance $D_A(z)$ (Mpc/h)')
plt.legend(loc='upper left')
plt.show()

"""
Explanation: Compare the distance functions. Line styles are solid for $D_C(z)$ and dashed for $D_A(z)$.
End of explanation
"""


izmin = 10
print 'Shear power calculation truncated at z >= %.2f' % z[izmin]

DA_min = min(flat_DA[izmin], open_DA[izmin], closed_DA[izmin])
DA_max = max(flat_DA[-1], open_DA[-1], closed_DA[-1])
print 'Using %.1f Mpc/h < DA < %.1f Mpc/h' % (DA_min, DA_max)

k_min, k_max = ell[0] / DA_max, ell[-1] / DA_min
print 'Using %.4f h/Mpc <= k <= %.4f h/Mpc' % (k_min, k_max)

"""
Explanation: Define the range of 3D wavenumbers $k = \ell/D_A$ needed to cover the calculation of the shear power. We have to set a minimum value of $D_A$ to use in order to establish a finite upper limit on $k$. Units are Mpc/h.
End of explanation
"""


flat_power = calculate_power(flat_model, k_min=k_min, k_max=k_max, scaled_by_h=True)
open_power = calculate_power(open_model, k_min=k_min, k_max=k_max, scaled_by_h=True)
closed_power = calculate_power(closed_model, k_min=k_min, k_max=k_max, scaled_by_h=True)

plt.plot(open_power['k'], open_power['Pk'], label='open')
plt.plot(flat_power['k'], flat_power['Pk'], label='flat')
plt.plot(closed_power['k'], closed_power['Pk'], label='closed')
plt.legend(loc='upper right')
plt.yscale('log')
plt.xscale('log')
plt.xlabel('3D wavenumber $k$ (h/Mpc)')
plt.ylabel('Linear matter power $P(k)$ (Mpc/h)$^3$')
plt.grid()
plt.show()

"""
Explanation: Calculate the matter power at z=0 for each of the fiducial cosmologies using the optional CLASS package. Units are h/Mpc for $k$ and (Mpc/h)**3 for $P(k)$.
End of explanation
"""


flat_growth = get_growth_function(flat_model, z)
open_growth = get_growth_function(open_model, z)
closed_growth = get_growth_function(closed_model, z)

plt.plot(z, open_growth, label='open')
plt.plot(z, flat_growth, label='flat')
plt.plot(z, closed_growth, label='closed')
plt.legend(loc='upper right')
plt.xlabel('Redshift $z$')
plt.ylabel('Growth function $G(z)$')
plt.grid()
plt.show()

"""
Explanation: Calculate the growth function for each cosmology:
End of explanation
"""


flat_variances = tabulate_3D_variances(ell, flat_DA[izmin:], flat_growth[izmin:], flat_power)
open_variances = tabulate_3D_variances(ell, open_DA[izmin:], open_growth[izmin:], open_power)
closed_variances = tabulate_3D_variances(ell, closed_DA[izmin:], closed_growth[izmin:], closed_power)

colors = ('r', 'g', 'b')
for j,iell in enumerate(range(0, 101, 50)):
    plt.plot(z[izmin:], flat_variances[iell], color=colors[j], label='$\ell = %.0f$' % ell[iell])
    plt.plot(z[izmin:], open_variances[iell], color=colors[j], ls='--')
    plt.plot(z[izmin:], closed_variances[iell], color=colors[j], ls=':')
plt.legend(loc='upper right')
plt.yscale('log')
plt.xlabel('Lens redshift $z_{src}$')
plt.ylabel('3D variance $(\pi/\ell) \Delta^2_{\delta}(k = \ell/D_A, z_{lens})$')
plt.grid()
plt.show()

"""
Explanation: Calculate the 3D variance contributions $V(\ell, D_A)$ for each cosmology, with $D_A > D_A(z_{min})$.
End of explanation
"""


flat_shear_power = calculate_shear_power(flat_DC[izmin:], flat_DA[izmin:],
                                         flat_weights[izmin:,izmin:], flat_variances)
open_shear_power = calculate_shear_power(open_DC[izmin:], open_DA[izmin:],
                                         open_weights[izmin:,izmin:], open_variances)
closed_shear_power = calculate_shear_power(closed_DC[izmin:], closed_DA[izmin:],
                                           closed_weights[izmin:,izmin:], closed_variances)

"""
Explanation: Shear-Shear Auto Power
Calculate the shear power $\Delta^2_{EE}(z_{src}, \ell)$ as a function of source position and 2D wavenumber $\ell$.
End of explanation
"""


def plot_shear_power():
    colors = ('r', 'g', 'b', 'y', 'k')
    for j,iz in enumerate(range(50, 251, 50)):
        plt.plot(ell, flat_shear_power[iz - izmin], color=colors[j], label='$z_{src}$ = %.2f' % z[iz])
        plt.plot(ell, open_shear_power[iz - izmin], color=colors[j], ls='--')
        plt.plot(ell, closed_shear_power[iz - izmin], color=colors[j], ls=':')
    plt.legend(loc='upper left')
    plt.yscale('log')
    plt.xscale('log')
    plt.xlabel('2D wavenumber $\ell$')
    plt.ylabel('Lensing shear power $\Delta^2_{EE}(z_{src}, \ell)$')
    plt.xlim(ell[0],ell[-1])
    y_max = max(np.max(flat_shear_power), np.max(open_shear_power), np.max(closed_shear_power))
    plt.ylim(None, 1.5 * y_max)
    plt.grid()
    #plt.savefig('shearpower.png')
    plt.show()

plot_shear_power()

"""
Explanation: Compare the calculated shear powers. Line styles are solid for the flat universe, dashed for the open universe, and dotted for the closed universe.
End of explanation
"""


flat_EE_cross = calculate_shear_power(flat_DC[izmin:], flat_DA[izmin:],
                                      flat_weights[izmin:,izmin:], flat_variances, 
                                      mode='shear-shear-cross')
flat_Eg_cross = calculate_shear_power(flat_DC[izmin:], flat_DA[izmin:],
                                      flat_weights[izmin:,izmin:], flat_variances, 
                                      mode='shear-galaxy-cross')

def plot_vs_z1_z2(data, exponent=0, label='f(z_1,z_2)'):
    scale = 10**exponent
    if np.max(data) == 0:
        cmap = 'CMRmap'
    else:
        cmap = 'CMRmap_r'
    plt.pcolormesh(z[izmin:], z[izmin:], scale * data, cmap=cmap)
    plt.colorbar(pad=0.05)
    cs = plt.contour(z[izmin:], z[izmin:], scale * data, colors='w')
    plt.clabel(cs, fmt='%.2f')
    plt.xlim(z[izmin],z[-1])
    plt.ylim(z[izmin],z[-1])
    plt.gca().set_aspect(1)
    plt.grid()
    plt.ylabel('Galaxy redshift $z_1$')
    plt.xlabel('Galaxy redshift $z_2$')
    if exponent != 0:
        label = ('$10^{%d} ' % exponent) + label
    plt.annotate(label, xy=(0.5, 0.01), xytext=(0.5, 0.01), color='k',
                 xycoords='axes fraction', textcoords='axes fraction',
                 horizontalalignment='center', verticalalignment='bottom',
                 fontsize='large', fontweight='bold')

plt.figure(figsize=(12,6.5))
# Plot EE cross spectra on the top row for ell = 10, 100, 1000.
plt.subplot(2,3,1)
plot_vs_z1_z2(flat_EE_cross[:,:,0], exponent=6, label='\Delta^2_{EE}(z_1,z_2,\ell=%.0f)$' % ell[0])
plt.subplot(2,3,2)
plot_vs_z1_z2(flat_EE_cross[:,:,50], exponent=6, label='\Delta^2_{EE}(z_1,z_2,\ell=%.0f)$' % ell[50])
plt.subplot(2,3,3)
plot_vs_z1_z2(flat_EE_cross[:,:,100], exponent=6, label='\Delta^2_{EE}(z_1,z_2,\ell=%.0f)$' % ell[100])
# Plot Eg cross spectra on the bottom row for ell = 10, 100, 1000.
plt.subplot(2,3,4)
plot_vs_z1_z2(flat_Eg_cross[:,:,0], exponent=2, label='\Delta^2_{Eg}(z_1,z_2,\ell=%.0f)/b_g$' % ell[0])
plt.subplot(2,3,5)
plot_vs_z1_z2(flat_Eg_cross[:,:,50], exponent=2, label='\Delta^2_{Eg}(z_1,z_2,\ell=%.0f)/b_g$' % ell[50])
plt.subplot(2,3,6)
plot_vs_z1_z2(flat_Eg_cross[:,:,100], exponent=2, label='\Delta^2_{Eg}(z_1,z_2,\ell=%.0f)/b_g$' % ell[100])
#
plt.tight_layout()
#plt.savefig('cross-power.png')
plt.show()

"""
Explanation: Shear-Shear and Shear-Galaxy Cross Power
Calculate the flat universe shear-shear cross power $\Delta^2_{EE}(z_1, z_2, \ell)$ and shear-galaxy cross power $\Delta^2_{Eg}(z_1, z_2, \ell)$ as a function of source positions $z_1$ and $z_2$, and 2D wavenumber $\ell$.
End of explanation
"""


theta_rad = 2 * np.pi / ell[::-1]
theta_arcmin = 60 * np.rad2deg(theta_rad)

"""
Explanation: Shear-Shear and Shear-Galaxy Cross-Correlation Functions
Define a log-spaced grid of angular separations where the correlation functions should be tabulated. Separations must be within the limits covered by our $\ell$ range, but do not need to be exactly $2\pi / \ell$.
End of explanation
"""


flat_EE_xi_p = calculate_correlation_function(flat_EE_cross, ell, theta_rad, order=0)
flat_EE_xi_m = calculate_correlation_function(flat_EE_cross, ell, theta_rad, order=4)
flat_Eg_xi = calculate_correlation_function(flat_Eg_cross, ell, theta_rad, order=2)

plt.figure(figsize=(12,9.75))
# Plot xi+ cross-correlations on the top row for 3 values of dtheta
plt.subplot(3,3,1)
plot_vs_z1_z2(flat_EE_xi_p[:,:,0], exponent=7, label="\\xi_+(z_1,z_2,\Delta\\theta=%.1f')$" % theta_arcmin[0])
plt.subplot(3,3,2)
plot_vs_z1_z2(flat_EE_xi_p[:,:,50], exponent=7, label="\\xi_+(z_1,z_2,\Delta\\theta=%.1f')$" % theta_arcmin[50])
plt.subplot(3,3,3)
plot_vs_z1_z2(flat_EE_xi_p[:,:,100], exponent=7, label="\\xi_+(z_1,z_2,\Delta\\theta=%.1f')$" % theta_arcmin[100])
# Plot xi- cross-correlations on the middle row for the same 3 values of dtheta
plt.subplot(3,3,4)
plot_vs_z1_z2(flat_EE_xi_m[:,:,0], exponent=7, label="\\xi_-(z_1,z_2,\Delta\\theta=%.1f')$" % theta_arcmin[0])
plt.subplot(3,3,5)
plot_vs_z1_z2(flat_EE_xi_m[:,:,50], exponent=7, label="\\xi_-(z_1,z_2,\Delta\\theta=%.1f')$" % theta_arcmin[50])
plt.subplot(3,3,6)
plot_vs_z1_z2(flat_EE_xi_m[:,:,100], exponent=7, label="\\xi_-(z_1,z_2,\Delta\\theta=%.1f')$" % theta_arcmin[100])
# Plot xi(Eg) cross-correlations on the bottom row for the same 3 values of dtheta
plt.subplot(3,3,7)
plot_vs_z1_z2(flat_Eg_xi[:,:,0], exponent=3, label="\\xi_{Eg}(z_1,z_2,\Delta\\theta=%.1f')/b_g$" % theta_arcmin[0])
plt.subplot(3,3,8)
plot_vs_z1_z2(flat_Eg_xi[:,:,50], exponent=3, label="\\xi_{Eg}(z_1,z_2,\Delta\\theta=%.1f')/b_g$" % theta_arcmin[50])
plt.subplot(3,3,9)
plot_vs_z1_z2(flat_Eg_xi[:,:,100], exponent=3, label="\\xi_{Eg}(z_1,z_2,\Delta\\theta=%.1f')/b_g$" % theta_arcmin[100])
#
plt.tight_layout()
#plt.savefig('cross-xi.png')
plt.show()

"""
Explanation: Calculate the shear-shear correlation functions, $\xi_+(z_1,z_2,\Delta\theta)$ and $\xi_-(z_1,z_2,\Delta\theta)$, and the shear-galaxy correlation function $\xi_{Eg}(z_1,z_2,\Delta\theta)$.
End of explanation
"""


bin1 = np.random.normal(0.55, 0.08, 50000)
bin1 = np.append(bin1, np.random.normal(0.35, 0.09, 50000))
bin2 = np.random.normal(0.70, 0.12, 100000)
bin3 = np.random.normal(0.90, 0.12, 50000)
bin3 = np.append(bin3, np.random.normal(1.20, 0.18, 50000))

plt.hist(bin1, bins=100, range=(0, 1.8), histtype='stepfilled', edgecolor='blue', facecolor='none');
plt.hist(bin2, bins=100, range=(0, 1.8), histtype='stepfilled', edgecolor='red', facecolor='none');
plt.hist(bin3, bins=100, range=(0, 1.8), histtype='stepfilled', edgecolor='green', facecolor='none');

"""
Explanation: Tomographic Predictions for DES Photo-z Bins
Approximately reproduce the DES photo-z redshift bins of Becker 2015 (Figure 3):
End of explanation
"""


def calculate_weights(data1, data2, points):
    num_points = len(points)
    # Histogram the datasets using the points as bin edges, resulting in num_points bin contents.
    data1_hist, edges = np.histogram(data1, points)
    data2_hist, edges = np.histogram(data2, points)
    # Calculate the joint pdf in num_points x num_points bins
    data1_pdf = data1_hist.astype(float) / data1.size
    data2_pdf = data2_hist.astype(float) / data2.size
    data12_pdf = data1_pdf[:, np.newaxis] * data2_pdf
    # Split each bin's probability equally between its four corner points.
    weights = np.zeros((num_points, num_points), dtype=float)
    weights[:-1, :-1] += data12_pdf
    weights[:-1, 1:] += data12_pdf
    weights[1:, :-1] += data12_pdf
    weights[1:, 1:] += data12_pdf
    weights /= 4
    return weights

w11 = calculate_weights(bin1, bin1, z[izmin:])
w21 = calculate_weights(bin2, bin1, z[izmin:])
w31 = calculate_weights(bin3, bin1, z[izmin:])
w22 = calculate_weights(bin2, bin2, z[izmin:])
w32 = calculate_weights(bin3, bin2, z[izmin:])
w33 = calculate_weights(bin3, bin3, z[izmin:])

plt.contour(z[izmin:], z[izmin:], w11, 2, colors='r')
plt.contour(z[izmin:], z[izmin:], w21, 2, colors='k')
plt.contour(z[izmin:], z[izmin:], w31, 2, colors='g')
plt.contour(z[izmin:], z[izmin:], w22, 2, colors='r')
plt.contour(z[izmin:], z[izmin:], w32, 2, colors='k')
plt.contour(z[izmin:], z[izmin:], w33, 2, colors='r')
plt.xlim(0, 1.5)
plt.ylim(0, 1.5)
plt.ylabel('Galaxy redshift $z_1$')
plt.xlabel('Galaxy redshift $z_2$')
plt.gca().set_aspect(1)

"""
Explanation: Calculate normalized weights for each pair of tomographic bins, on the same $(z_1, z_2)$ grid used to calculate cross spectra and cross correlations.
End of explanation
"""


flat_EE_xi_p_11 = np.sum(flat_EE_xi_p * w11[:, :, np.newaxis], axis=(0, 1))
flat_EE_xi_p_21 = np.sum(flat_EE_xi_p * w21[:, :, np.newaxis], axis=(0, 1))
flat_EE_xi_p_31 = np.sum(flat_EE_xi_p * w31[:, :, np.newaxis], axis=(0, 1))
flat_EE_xi_p_22 = np.sum(flat_EE_xi_p * w22[:, :, np.newaxis], axis=(0, 1))
flat_EE_xi_p_32 = np.sum(flat_EE_xi_p * w32[:, :, np.newaxis], axis=(0, 1))
flat_EE_xi_p_33 = np.sum(flat_EE_xi_p * w33[:, :, np.newaxis], axis=(0, 1))

flat_EE_xi_m_11 = np.sum(flat_EE_xi_m * w11[:, :, np.newaxis], axis=(0, 1))
flat_EE_xi_m_21 = np.sum(flat_EE_xi_m * w21[:, :, np.newaxis], axis=(0, 1))
flat_EE_xi_m_31 = np.sum(flat_EE_xi_m * w31[:, :, np.newaxis], axis=(0, 1))
flat_EE_xi_m_22 = np.sum(flat_EE_xi_m * w22[:, :, np.newaxis], axis=(0, 1))
flat_EE_xi_m_32 = np.sum(flat_EE_xi_m * w32[:, :, np.newaxis], axis=(0, 1))
flat_EE_xi_m_33 = np.sum(flat_EE_xi_m * w33[:, :, np.newaxis], axis=(0, 1))

def plot_xi(axes, row, col, xi, label):
    axis = axes[row, col]
    axis.plot(theta_arcmin, 1e4 * theta_arcmin * xi)
    axis.set_xscale('log')
    if row == len(axes) - 1:
        axis.set_xlabel('$\\theta$ [arcmin]')
    if col == 0:
        axis.set_ylabel('$\\theta \\times \\xi_%s(\\theta) / 10^{-4}$' % ('-' if col%2 else '-'))
    axis.set_xlim(2., 300.)
    axis.set_ylim(-4.25, +8.5)
    axis.grid(True)
    axis.annotate(label, xy=(0.1,0.9), xytext=(0.1,0.9),
                  xycoords='axes fraction', textcoords='axes fraction',
                  horizontalalignment='left', verticalalignment='top',
                  fontsize='large', fontweight='bold')

"""
Explanation: Integrate the $\xi_\pm(\theta)$ correlation functions over the PDF of each pair of bins:
End of explanation
"""


fig,axes = plt.subplots(6, 3, sharex='col', sharey='row', figsize=(12, 12))
fig.subplots_adjust(hspace=0, wspace=0)
plot_xi(axes, 0, 0, flat_EE_xi_p_33, '3-3')
plot_xi(axes, 1, 0, flat_EE_xi_m_33, '3-3')
plot_xi(axes, 2, 0, flat_EE_xi_p_32, '3-2')
plot_xi(axes, 3, 0, flat_EE_xi_m_32, '3-2')
plot_xi(axes, 2, 1, flat_EE_xi_p_22, '2-2')
plot_xi(axes, 3, 1, flat_EE_xi_m_22, '2-2')
plot_xi(axes, 4, 0, flat_EE_xi_p_31, '3-1')
plot_xi(axes, 5, 0, flat_EE_xi_m_31, '3-1')
plot_xi(axes, 4, 1, flat_EE_xi_p_21, '2-1')
plot_xi(axes, 5, 1, flat_EE_xi_m_21, '2-1')
plot_xi(axes, 4, 2, flat_EE_xi_p_11, '1-1')
plot_xi(axes, 5, 2, flat_EE_xi_m_11, '1-1')
#plt.savefig('DES-fig2.png')
plt.show()

"""
Explanation: Try to approximately reproduce Fig.2 of Becker 2015 (including the large y-axis range, which makes precise comparisons difficult).  Note that our calculation only goes down to about 20 arcsecs, so does not cover the full range of the DES plot.
End of explanation
"""


