halo_comparison/spectra_computation.py

# Call with spectra_computation.py <time> <kind>
# time = 'ics' for ICs, = 'z=1' for redshift z=1, = 'end' for final results
# kind = 'power' for power spectra comparing same resolution, 'cross' for comparing across all resolutions

import itertools
import subprocess
from pathlib import Path
from sys import argv

from jobrun.jobrun import jobrun
from paths import base_dir, account, qos, spectra_dir
from spectra_plot import waveforms

vsc = True


def spectra_jobrun(args, maxres: int):
    mem = 380 if maxres == 1048 else 280
    if vsc:
        jobrun(
            args, time="12:00:00", nodes=1, mem=mem,
            source=Path("/gpfs/data/fs71636/lwinkler/spack-latest.sh"),
            qos=qos, account=account
        )
    else:
        subprocess.run(args, check=True)


def run_spectra(
        waveform: str, resolution_1: int, resolution_2: int, Lbox: int, time: str
):
    print("starting")
    setup_1 = f"{waveform}_{resolution_1}_{Lbox}"
    setup_2 = f"{waveform}_{resolution_2}_{Lbox}"

    # #For ICs: time == 'ics'
    if time == "ics":
        output_file = (
                base_dir
                / f"spectra/{waveform}_{Lbox}/{waveform}_{Lbox}_ics_{resolution_1}_{resolution_2}_cross_spectrum.txt"
        )
        output_file.parent.mkdir(exist_ok=True, parents=True)
        if output_file.exists():
            print(f"{output_file} already exists, skipping.")
            return
        spectra_jobrun(
            [
                str(spectra),
                "--ngrid",
                "2048",
                "--format=4",  # This seems to work, but is not as readable
                "--output",
                str(
                    base_dir
                    / f"spectra/{waveform}_{Lbox}/{waveform}_{Lbox}_ics_{resolution_1}_{resolution_2}"
                ),
                "--input",
                str(base_dir / f"{setup_1}/ics_{setup_1}.hdf5"),
                "--input",
                str(base_dir / f"{setup_2}/ics_{setup_2}.hdf5"),
            ], max(resolution_1, resolution_2)
        )

    # #For evaluation of results at redshift z=1: time == 'z=1'
    elif time == "z=1":
        output_file = (
                base_dir
                / f"spectra/{waveform}_{Lbox}/{waveform}_{Lbox}_a2_{resolution_1}_{resolution_2}_cross_spectrum.txt"
        )
        output_file.parent.mkdir(exist_ok=True, parents=True)
        if output_file.exists():
            print(f"{output_file} already exists, skipping.")
            return
        spectra_jobrun(
            [
                str(spectra),
                "--ngrid",
                "2048",
                "--format=3",
                "--output",
                str(
                    base_dir
                    / f"spectra/{waveform}_{Lbox}/{waveform}_{Lbox}_a2_{resolution_1}_{resolution_2}"
                ),
                "--input",
                str(base_dir / f"{setup_1}/output_0002.hdf5"),
                "--input",
                str(base_dir / f"{setup_2}/output_0002.hdf5"),
            ], max(resolution_1, resolution_2)
        )


    # #For evaluation of final results: time == 'end'
    elif time == "end":
        output_file = (
                base_dir
                / f"spectra/{waveform}_{Lbox}/{waveform}_{Lbox}_a4_{resolution_1}_{resolution_2}_cross_spectrum.txt"
        )
        output_file.parent.mkdir(exist_ok=True, parents=True)
        if output_file.exists():
            print(f"{output_file} already exists, skipping.")
            return
        spectra_jobrun(
            [
                str(spectra),
                "--ngrid",
                "2048",
                "--format=3",
                "--output",
                str(
                    base_dir
                    / f"spectra/{waveform}_{Lbox}/{waveform}_{Lbox}_a4_{resolution_1}_{resolution_2}"
                ),
                "--input",
                str(base_dir / f"{setup_1}/output_0004.hdf5"),
                "--input",
                str(base_dir / f"{setup_2}/output_0004.hdf5"),
            ], max(resolution_1, resolution_2)
        )
    else:
        raise ValueError(f"invalid time ({time})")

    print("end")


def power_run(resolutions: list, Lbox: int, time: str):
    args = []
    for waveform in waveforms:
        for resolution in resolutions:
            args.append((waveform, resolution, resolution, Lbox, time))
    return args


def cross_run(resolutions: list, Lbox: int, time: str):
    args = []
    for waveform in waveforms:
        for res1, res2 in itertools.combinations(resolutions, 2):
            args.append((waveform, res1, res2, Lbox, time))
    return args


if __name__ == "__main__":
    #    input("are you sure you want to run this? This might need a large amount of memory")
    Lbox = 100
    resolutions = [128, 256, 512, 1024]

    spectra = spectra_dir / "spectra"
    time = argv[1]

    if argv[2] == "power":
        args = power_run(resolutions=resolutions, Lbox=Lbox, time=time)

    elif argv[2] == "cross":
        args = cross_run(resolutions=resolutions, Lbox=Lbox, time=time)
    else:
        raise ValueError("missing argv[2] (power|cross)")
    for a in args:
        run_spectra(*a)
    # with Pool(processes=1) as p:
    #     p.starmap(run_spectra, args)