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 multiprocessing import Pool, cpu_count
from sys import argv

from paths import base_dir, spectra_dir
from spectra_plot import waveforms


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'
        if output_file.exists():
            print(f'{output_file} already exists, skipping.')
            return
        subprocess.run([
            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')
        ], check=True)
    
    # #For evaluation of results at redshift z=1: time == 'z=1' | NOT ADAPTED FOR VSC5 YET!
    elif time == 'z=1':
        output_file = base_dir / f'spectra/{waveform}_{Lbox}/{waveform}_{Lbox}_a2_{resolution_1}_{resolution_2}_cross_spectrum.txt'
        if output_file.exists():
            print(f'{output_file} already exists, skipping.')
            return
        subprocess.run([
            str(spectra),
            '--ngrid',
            '1024',
            '--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')
        ], check=True)

    # #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'
        if output_file.exists():
            print(f'{output_file} already exists, skipping.')
            return
        subprocess.run([
            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')
        ], check=True)
    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)")
    with Pool(processes=1) as p:
        p.starmap(run_spectra, args)