cmb_modules.py

import numpy as np
import matplotlib
import sys
import matplotlib.cm as cm
import matplotlib.mlab as mlab
import matplotlib.pyplot as plt
import astropy.io.fits as fits


def make_CMB_T_map(N, pix_size, ell, DlTT):
    "makes a realization of a simulated CMB sky map given an input DlTT as a function of ell,"
    "the pixel size (pix_size) required and the number N of pixels in the linear dimension."
    # np.random.seed(100)
    # convert Dl to Cl
    ClTT = DlTT * 2 * np.pi / (ell * (ell + 1.0))
    ClTT[0] = 0.0  # set the monopole and the dipole of the Cl spectrum to zero
    ClTT[1] = 0.0

    # make a 2D real space coordinate system
    inds = np.linspace(-0.5, 0.5, N)
    X, Y = np.meshgrid(inds, inds)
    # radial component R
    R = np.sqrt(X**2.0 + Y**2.0)

    # now make a 2D CMB power spectrum
    pix_to_rad = (
        pix_size / 60.0 * np.pi / 180.0
    )  # going from pix_size in arcmins to degrees and then degrees to radians
    ell_scale_factor = (
        2.0 * np.pi / pix_to_rad
    )  # now relating the angular size in radians to multipoles
    ell2d = (
        R * ell_scale_factor
    )  # making a fourier space analogue to the real space R vector
    ClTT_expanded = np.zeros(int(ell2d.max()) + 1)
    # making an expanded Cl spectrum (of zeros) that goes all the way to the size of the 2D ell vector
    ClTT_expanded[
        0 : (ClTT.size)
    ] = ClTT  # fill in the Cls until the max of the ClTT vector

    # the 2D Cl spectrum is defined on the multiple vector set by the pixel scale
    CLTT2d = ClTT_expanded[ell2d.astype(int)]
    # plt.imshow(np.log(CLTT2d))

    # now make a realization of the CMB with the given power spectrum in real space
    random_array_for_T = np.random.normal(0, 1, (N, N))
    FT_random_array_for_T = np.fft.fft2(
        random_array_for_T
    )  # take FFT since we are in Fourier space

    FT_2d = (
        np.sqrt(CLTT2d) * FT_random_array_for_T
    )  # we take the sqrt since the power spectrum is T^2
    # plt.imshow(np.real(FT_2d))

    ## make a plot of the 2D cmb simulated map in Fourier space, note the x and y axis labels need to be fixed
    # Plot_CMB_Map(np.real(np.conj(FT_2d)*FT_2d*ell2d * (ell2d+1)/2/np.pi),0,np.max(np.conj(FT_2d)*FT_2d*ell2d * (ell2d+1)/2/np.pi),ell2d.max(),ell2d.max())  ###

    # move back from ell space to real space
    CMB_T = np.fft.ifft2(np.fft.fftshift(FT_2d))
    # move back to pixel space for the map
    CMB_T = CMB_T / pix_to_rad
    # we only want to plot the real component
    CMB_T = np.real(CMB_T)

    ## return the map
    return CMB_T


def Plot_CMB_Map(Map_to_Plot, c_min, c_max, X_width, Y_width):
    from mpl_toolkits.axes_grid1 import make_axes_locatable

    print("map mean:", np.mean(Map_to_Plot), "map rms:", np.std(Map_to_Plot))
    plt.gcf().set_size_inches(10, 10)
    im = plt.imshow(
        Map_to_Plot, interpolation="bilinear", origin="lower", cmap=cm.RdBu_r
    )
    im.set_clim(c_min, c_max)
    im.set_extent([0, X_width, 0, Y_width])
    plt.ylabel(r"angle $[^\circ]$")
    plt.xlabel(r"angle $[^\circ]$")

    ax = plt.gca()
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    cbar = plt.colorbar(im, cax=cax, label="Temperature [uK]")

    plt.show()


def Poisson_source_component(N, pix_size, Number_of_Sources, Amplitude_of_Sources):
    "makes a realization of a naive Poisson-distributed point source map"
    PSMap = np.zeros([int(N), int(N)])
    i = 0
    print("Number of sources required: ", Number_of_Sources)

    while i < int(Number_of_Sources):
        pix_x = int(N * np.random.rand())
        pix_y = int(N * np.random.rand())
        PSMap[pix_x, pix_y] += np.random.poisson(Amplitude_of_Sources)
        i = i + 1

    return PSMap


###############################


def Exponential_source_component(
    N, pix_size, Number_of_Sources_EX, Amplitude_of_Sources_EX
):
    N = int(N)
    "makes a realization of a naive exponentially-distributed point source map"
    PSMap = np.zeros([N, N])
    i = 0
    while i < Number_of_Sources_EX:
        pix_x = int(N * np.random.rand())
        pix_y = int(N * np.random.rand())
        PSMap[pix_x, pix_y] += np.random.exponential(Amplitude_of_Sources_EX)
        i = i + 1

    return PSMap


###############################


def SZ_source_component(
    N,
    pix_size,
    Number_of_SZ_Clusters,
    Mean_Amplitude_of_SZ_Clusters,
    SZ_beta,
    SZ_Theta_core,
    do_plots,
):
    "makes a realization of a nieve SZ map"
    N = int(N)
    SZMap = np.zeros([N, N])
    SZcat = np.zeros(
        [3, Number_of_SZ_Clusters]
    )  ## catalogue of SZ sources, X, Y, amplitude
    # make a distribution of point sources with varying amplitude
    i = 0
    while i < Number_of_SZ_Clusters:
        pix_x = int(N * np.random.rand())
        pix_y = int(N * np.random.rand())
        pix_amplitude = np.random.exponential(Mean_Amplitude_of_SZ_Clusters) * (-1.0)
        SZcat[0, i] = pix_x
        SZcat[1, i] = pix_y
        SZcat[2, i] = pix_amplitude
        SZMap[pix_x, pix_y] += pix_amplitude
        i = i + 1
    if do_plots:
        hist, bin_edges = np.histogram(SZMap, bins=50, range=[SZMap.min(), -10])
        plt.semilogy(bin_edges[0:-1], hist)
        plt.xlabel(r"source amplitude [$\mu$K]")
        plt.ylabel("number or pixels")
        plt.show()

    # make a beta function
    beta = beta_function(N, pix_size, SZ_beta, SZ_Theta_core)

    # convovle the beta funciton with the point source amplitude to get the SZ map
    FT_beta = np.fft.fft2(np.fft.fftshift(beta))
    FT_SZMap = np.fft.fft2(np.fft.fftshift(SZMap))
    SZMap = np.fft.fftshift(np.real(np.fft.ifft2(FT_beta * FT_SZMap)))

    # return the SZ map
    return (SZMap, SZcat)


###############################


def beta_function(N, pix_size, SZ_beta, SZ_Theta_core):
    # make a beta function
    N = int(N)
    ones = np.ones(N)
    inds = (np.arange(N) + 0.5 - N / 2.0) * pix_size
    X = np.outer(ones, inds)
    Y = np.transpose(X)
    R = np.sqrt(X**2.0 + Y**2.0)

    beta = (1 + (R / SZ_Theta_core) ** 2.0) ** ((1 - 3.0 * SZ_beta) / 2.0)

    # return the beta function map
    return beta


###############################


def convolve_map_with_gaussian_beam(N, pix_size, beam_size_fwhp, Map):
    "convolves a map with a gaussian beam pattern.  NOTE: pix_size and beam_size_fwhp need to be in the same units"
    # make a 2d gaussian
    gaussian = make_2d_gaussian_beam(N, pix_size, beam_size_fwhp)

    # do the convolution
    FT_gaussian = np.fft.fft2(np.fft.fftshift(gaussian))
    FT_Map = np.fft.fft2(np.fft.fftshift(Map))
    convolved_map = np.fft.fftshift(np.real(np.fft.ifft2(FT_gaussian * FT_Map)))

    # return the convolved map
    return convolved_map


###############################


def make_2d_gaussian_beam(N, pix_size, beam_size_fwhp):
    # make a 2d coordinate system
    N = int(N)
    ones = np.ones(N)
    inds = (np.arange(N) + 0.5 - N / 2.0) * pix_size
    X = np.outer(ones, inds)
    Y = np.transpose(X)
    R = np.sqrt(X**2.0 + Y**2.0)

    # make a 2d gaussian
    beam_sigma = beam_size_fwhp / np.sqrt(8.0 * np.log(2))
    gaussian = np.exp(-0.5 * (R / beam_sigma) ** 2.0)
    gaussian = gaussian / np.sum(gaussian)

    # return the gaussian
    return gaussian


###############################


def make_noise_map(
    N, pix_size, white_noise_level, atmospheric_noise_level, one_over_f_noise_level
):
    "makes a realization of instrument noise, atmosphere and 1/f noise level set at 1 degrees"
    ## make a white noise map
    N = int(N)
    white_noise = np.random.normal(0, 1, (N, N)) * white_noise_level / pix_size

    ## make an atmosperhic noise map
    atmospheric_noise = 0.0
    if atmospheric_noise_level != 0:
        ones = np.ones(N)
        inds = np.arange(N) + 0.5 - N / 2.0
        X = np.outer(ones, inds)
        Y = np.transpose(X)
        R = (
            np.sqrt(X**2.0 + Y**2.0) * pix_size / 60.0
        )  ## angles relative to 1 degrees
        mag_k = 2 * np.pi / (R + 0.01)  ## 0.01 is a regularizaiton factor
        atmospheric_noise = np.fft.fft2(np.random.normal(0, 1, (N, N)))
        atmospheric_noise = (
            np.fft.ifft2(atmospheric_noise * np.fft.fftshift(mag_k ** (5 / 3.0)))
            * atmospheric_noise_level
            / pix_size
        )

    ## make a 1/f map, along a single direction to illustrate striping
    oneoverf_noise = 0.0
    if one_over_f_noise_level != 0:
        ones = np.ones(N)
        inds = np.arange(N) + 0.5 - N / 2.0
        X = np.outer(ones, inds) * pix_size / 60.0  ## angles relative to 1 degrees
        kx = 2 * np.pi / (X + 0.01)  ## 0.01 is a regularizaiton factor
        oneoverf_noise = np.fft.fft2(np.random.normal(0, 1, (N, N)))
        oneoverf_noise = (
            np.fft.ifft2(oneoverf_noise * np.fft.fftshift(kx))
            * one_over_f_noise_level
            / pix_size
        )

    ## return the noise map
    noise_map = np.real(white_noise + atmospheric_noise + oneoverf_noise)
    return noise_map


###############################
def Filter_Map(Map, N, N_mask):
    N = int(N)
    ## set up a x, y, and r coordinates for mask generation
    ones = np.ones(N)
    inds = np.arange(N) + 0.5 - N / 2.0
    X = np.outer(ones, inds)
    Y = np.transpose(X)
    R = np.sqrt(X**2.0 + Y**2.0)  ## angles relative to 1 degrees

    ## make a mask
    mask = np.ones([N, N])
    mask[np.where(np.abs(X) < N_mask)] = 0

    return apply_filter(Map, mask)


def apply_filter(Map, filter2d):
    ## apply the filter in fourier space
    FMap = np.fft.fftshift(np.fft.ifft2(np.fft.fftshift(Map)))
    FMap_filtered = FMap * filter2d
    Map_filtered = np.real(np.fft.fftshift(np.fft.fft2(FMap_filtered)))

    ## return the output
    return Map_filtered


def cosine_window(N):
    "makes a cosine window for apodizing to avoid edges effects in the 2d FFT"
    # make a 2d coordinate system
    ones = np.ones(N)
    inds = (np.arange(N) + 0.5 - N / 2.0) / N * np.pi  ## eg runs from -pi/2 to pi/2
    X = np.outer(ones, inds)
    Y = np.transpose(X)

    # make a window map
    window_map = np.cos(X) * np.cos(Y)

    # return the window map
    return window_map


###############################


def average_N_spectra(spectra, N_spectra, N_ells):
    avgSpectra = np.zeros(N_ells)
    rmsSpectra = np.zeros(N_ells)

    # calcuate the average spectrum
    i = 0
    while i < N_spectra:
        avgSpectra = avgSpectra + spectra[i, :]
        i = i + 1
    avgSpectra = avgSpectra / (1.0 * N_spectra)

    # calculate the rms of the spectrum
    i = 0
    while i < N_spectra:
        rmsSpectra = rmsSpectra + (spectra[i, :] - avgSpectra) ** 2
        i = i + 1
    rmsSpectra = np.sqrt(rmsSpectra / (1.0 * N_spectra))

    return (avgSpectra, rmsSpectra)


def calculate_2d_spectrum(Map1, Map2, delta_ell, ell_max, pix_size, N):
    "calcualtes the power spectrum of a 2d map by FFTing, squaring, and azimuthally averaging"
    N = int(N)
    # make a 2d ell coordinate system
    ones = np.ones(N)
    inds = (np.arange(N) + 0.5 - N / 2.0) / (N - 1.0)
    kX = np.outer(ones, inds) / (pix_size / 60.0 * np.pi / 180.0)
    kY = np.transpose(kX)
    K = np.sqrt(kX**2.0 + kY**2.0)
    ell_scale_factor = 2.0 * np.pi
    ell2d = K * ell_scale_factor

    # make an array to hold the power spectrum results
    N_bins = int(ell_max / delta_ell)
    ell_array = np.arange(N_bins)
    CL_array = np.zeros(N_bins)

    # get the 2d fourier transform of the map
    FMap1 = np.fft.ifft2(np.fft.fftshift(Map1))
    FMap2 = np.fft.ifft2(np.fft.fftshift(Map2))
    PSMap = np.fft.fftshift(np.real(np.conj(FMap1) * FMap2))
    # fill out the spectra
    i = 0
    while i < N_bins:
        ell_array[i] = (i + 0.5) * delta_ell
        inds_in_bin = (
            (ell2d >= (i * delta_ell)) * (ell2d < ((i + 1) * delta_ell))
        ).nonzero()
        CL_array[i] = np.mean(PSMap[inds_in_bin])
        # print i, ell_array[i], inds_in_bin, CL_array[i]
        i = i + 1

    # return the power spectrum and ell bins
    return (ell_array, CL_array * np.sqrt(pix_size / 60.0 * np.pi / 180.0) * 2.0)


def Plot_CMB_Lensing_Map(Map_to_Plot, X_width, Y_width):
    c_max = np.max(Map_to_Plot)
    c_min = np.min(Map_to_Plot)
    from mpl_toolkits.axes_grid1 import make_axes_locatable

    print("map mean:", np.mean(Map_to_Plot), "map rms:", np.std(Map_to_Plot))
    plt.gcf().set_size_inches(10, 10)
    im = plt.imshow(Map_to_Plot, interpolation="bilinear", origin="lower", cmap="gray")
    im.set_clim(c_min, c_max)
    im.set_extent([0, X_width, 0, Y_width])
    plt.ylabel(r"angle $[^\circ]$")
    plt.xlabel(r"angle $[^\circ]$")

    ax = plt.gca()
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    cbar = plt.colorbar(im, cax=cax, label="Kappa [arb]")

    plt.show()