halo_comparison/spectra_plot.py

import itertools
from pathlib import Path
from sys import argv

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.axes import Axes
from matplotlib.figure import Figure

from paths import base_dir
from utils import figsize_from_page_fraction

Lbox = 100
h = 0.690021
k0 = 3.14159265358979323846264338327950 / Lbox
resolutions = [128, 256, 512, 1024]
waveforms = ["DB2", "DB4", "DB8", "shannon"]

# Careful: k is actually in Mpc^-1, the column is just named weirdly.
columns = [
    "k [Mpc]",
    "Pcross",
    "P1",
    "err. P1",
    "P2",
    "err. P2",
    "P2-1",
    "err. P2-1",
    "modes in bin",
]

# linestyles = ["solid", "dashed", "dotted"]
colors = [f"C{i}" for i in range(10)]


# colors = ["C1", "C2", "C3", "C4"]


def spectra_data(
        waveform: str, resolution_1: int, resolution_2: int, Lbox: int, time: str
):
    dir = base_dir / f"spectra/{waveform}_{Lbox}"

    if time == "ics":
        spectra_data = pd.read_csv(
            f"{dir}/{waveform}_{Lbox}_ics_{resolution_1}_{resolution_2}_cross_spectrum.txt",
            sep=" ",
            skipinitialspace=True,
            header=None,
            names=columns,
            skiprows=1,
        )
    elif time == "z=1":
        spectra_data = pd.read_csv(
            f"{dir}/{waveform}_{Lbox}_a2_{resolution_1}_{resolution_2}_cross_spectrum.txt",
            sep=" ",
            skipinitialspace=True,
            header=None,
            names=columns,
            skiprows=1,
        )
    elif time == "end":
        spectra_data = pd.read_csv(
            f"{dir}/{waveform}_{Lbox}_a4_{resolution_1}_{resolution_2}_cross_spectrum.txt",
            sep=" ",
            skipinitialspace=True,
            header=None,
            names=columns,
            skiprows=1,
        )
    else:
        raise ValueError(f"invalid time ({time}) should be (ics|z=1|end)")

    # only consider rows above resolution limit
    spectra_data = spectra_data[spectra_data["k [Mpc]"] >= k0]

    return spectra_data


def create_plot(mode):
    fig: Figure
    combination_list = list(itertools.combinations(resolutions, 2))
    fig, axes = plt.subplots(
        len(waveforms), 3, sharex=True, sharey=True,
        constrained_layout=True, figsize=figsize_from_page_fraction(columns=2),
    )
    crossings = np.zeros((len(waveforms), len(combination_list)))
    for i, waveform in enumerate(waveforms):
        ax_ics: Axes = axes[i][0]
        ax_z1: Axes = axes[i][1]
        ax_end: Axes = axes[i][2]
        axes_names = {
            # TODO: better names
            ax_ics: "ics",
            ax_z1: "z=1",
            ax_end: "end"
        }
        bottom_row = i == len(waveforms) - 1
        # for is_end, ax in enumerate([ax_ics, ax_z1]):
        for is_end, ax in enumerate([ax_ics, ax_z1, ax_end]):
            if bottom_row:
                ax.set_xlabel("k [Mpc$^{-1}$]")
            ax.text(
                0.02,
                0.85,
                f"{waveform}",
                horizontalalignment="left",
                verticalalignment="top",
                transform=ax.transAxes,
            )
            ax.text(
                0.98,
                0.85,
                axes_names[ax],
                horizontalalignment="right",
                verticalalignment="top",
                transform=ax.transAxes,
            )
            for j, res in enumerate(resolutions[:-1] if mode == "cross" else resolutions):
                ax.axvline(
                    k0 * res,
                    color=colors[j],
                    linestyle="dashed",
                    label=f"$k_\\mathrm{{ny, {res}}}$" if mode =="power" else None,
                )
            # ax.set_xticklabels([])
            # ax.set_yticklabels([])

        if mode == "power":
            ax_ics.set_ylabel("P")
            for j, resolution in enumerate(resolutions):
                ics_data = spectra_data(waveform, resolution, resolution, Lbox, "ics")
                ics_k = ics_data["k [Mpc]"]
                ics_p1 = ics_data["P1"]
                comp_data = spectra_data(waveform, resolutions[-1], resolutions[-1], Lbox, "ics")
                comp_p1 = comp_data["P1"]
                ics_p1 /= comp_p1

                end_data = spectra_data(waveform, resolution, resolution, Lbox, "end")
                end_k = end_data["k [Mpc]"]
                end_p1 = end_data["P1"]
                comp_data = spectra_data(waveform, resolutions[-1], resolutions[-1], Lbox, "end")
                comp_p1 = comp_data["P1"]
                end_p1 /= comp_p1

                z1_data = spectra_data(waveform, resolution, resolution, Lbox, "z=1")
                z1_k = z1_data["k [Mpc]"]
                z1_p1 = z1_data["P1"]
                comp_data = spectra_data(waveform, resolutions[-1], resolutions[-1], Lbox, 'z=1')
                comp_p1 = comp_data["P1"]
                z1_p1 /= comp_p1

                ax_ics.semilogx(ics_k, ics_p1, color=colors[j])
                ax_z1.semilogx(z1_k, z1_p1, color=colors[j])
                ax_end.semilogx(end_k, end_p1, color=colors[j])
                for ax in [ax_ics, ax_z1, ax_end]:
                    ax.set_ylim(0.9, 1.10)


        # fig.suptitle(f"Power Spectra {time}") #Not needed for paper
        # fig.tight_layout()

        elif mode == "cross":
            ax_ics.set_ylabel("C")
            # ax_end.set_ylabel("C")
            for j, (res1, res2) in enumerate(combination_list):
                ics_data = spectra_data(waveform, res1, res2, Lbox, 'ics')
                ics_k = ics_data["k [Mpc]"]
                ics_pcross = ics_data["Pcross"]
                smaller_res = min(res1, res2)
                crossing_index = np.searchsorted(ics_k.to_list(), k0 * smaller_res)
                crossing_value = ics_pcross[crossing_index]
                crossings[i][j] = crossing_value
                ax_ics.semilogx(ics_k, ics_pcross, color=colors[j + 3], label=f'{res1} vs {res2}')

                z1_data = spectra_data(waveform, res1, res2, Lbox, 'z=1')
                z1_k = z1_data["k [Mpc]"]
                z1_pcross = z1_data["Pcross"]

                ax_z1.semilogx(z1_k, z1_pcross, color=colors[j + 3], label=f'{res1} vs {res2}')

                end_data = spectra_data(waveform, res1, res2, Lbox, 'end')
                end_k = end_data["k [Mpc]"]
                end_pcross = end_data["Pcross"]

                ax_end.semilogx(end_k, end_pcross, color=colors[j + 3], label=f'{res1} vs {res2}')

            ax_end.set_xlim(right=k0 * resolutions[-1])
            ax_end.set_ylim(0.8, 1.02)
        if bottom_row:
            # ax_z1.legend()
            ax_ics.legend(loc='lower left')

        # fig.suptitle(f"Cross Spectra {time}") #Not needed for paper
        # fig.tight_layout()
    print(crossings)
    crossings_df = pd.DataFrame(crossings, columns=combination_list, index=waveforms)
    # print(crossings_df.to_markdown())
    print(crossings_df.to_latex())
    # fig.tight_layout()
    fig.savefig(Path(f"~/tmp/spectra_{mode}.pdf").expanduser())


def main():
    if len(argv) < 2:
        print("run spectra_plot.py [power|cross] or spectra_plot.py all")
        exit(1)
    if argv[1] == "all":
        for mode in ["power", "cross"]:
            create_plot(mode)
        plt.show()
        return
    mode = argv[1]
    create_plot(mode)
    plt.show()


if __name__ == "__main__":
    main()