swift-pipeline

#!/usr/bin/env python3
"""
SWIFT-pipeline is a significantly more complex version of ``velociraptor-plot``.

It uses configuration files along with additional plotting scripts to put together
webpages that can represent whole cosmological simulations easily.
"""

import argparse as ap
from typing import Tuple, List, Any
from p_tqdm import p_map
from time import time

parser = ap.ArgumentParser(
    prog="swift-pipeline",
    description=(
        "Creates a webpage containing many figures out of your SWIFT runs. "
        "When creating this page, also creates a metadata file that can be used "
        "later with this program to produce comparison webpages between multiple "
        "simulations."
    ),
    epilog=(
        "Example usage:\n"
        "swift-pipeline -C ~/config -c example_0000.properties -s snapshot_0000.hdf5 "
        "-o ~/plots/example_0000 -i /path/to/my/sim\n\n"
        "Example creating comparisons:\n"
        "swift-pipeline -C ~/config -c example_0000.properties example_0000.properties "
        "-s snapshot_0000.hdf5 snapshot_0000.hdf5 "
        "-o ~/plots/example_0000 -i /path/to/my/first/sim /path/to/my/second/sim\n\n"
    ),
)

parser.add_argument(
    "-C",
    "--config",
    type=str,
    required=True,
    help=("Configuration directory, containing config.yml."),
)

parser.add_argument(
    "-c",
    "--catalogues",
    type=str,
    required=True,
    help="Name of the VELOCIraptor HDF5 .properties file(s). Required.",
    nargs="*",
)

parser.add_argument(
    "-s",
    "--snapshots",
    required=True,
    type=str,
    help="Name of the snapshot file(s). Required.",
    nargs="*",
)

parser.add_argument(
    "-o",
    "--output",
    type=str,
    required=True,
    help="Output directory for figures. Required.",
)

parser.add_argument(
    "-i",
    "--input",
    type=str,
    required=False,
    default=".",
    help=(
        "Input directory where the snapshot(s) and properties file(s) are located. "
        "Default is the current working directory. If you are running for comparison "
        "purposes you will need to ensure that the metadata yaml files have been "
        "generated in these folders and have the same basename (--metadata) as is "
        "given here."
    ),
    nargs="*",
)

parser.add_argument(
    "-d",
    "--debug",
    required=False,
    default=False,
    action="store_true",
    help="Run in debug mode if this flag is present. Default: no.",
)


parser.add_argument(
    "-m",
    "--metadata",
    required=False,
    default="data",
    help=(
        "Base name of the written metadata file in the input directory. "
        "By default this is data, leading to data_XXXX.yml"
    ),
)

parser.add_argument(
    "-n",
    "--run-names",
    required=False,
    default=None,
    nargs="*",
    help=(
        "Overwrite the names given to each run? If not present, the default names "
        "from the snapshots are used, and in the case where there are multiple "
        "redshifts, we append the redshift."
    ),
)

parser.add_argument(
    "-j",
    "--num-of-cpus",
    required=False,
    type=int,
    default=None,
    help=(
        "Number of CPUs to use for running scripts in parallel. If not specified, uses "
        "the maximum number of CPUs avaliable in the system."
    ),
)

parser.add_argument(
    "-M",
    "--special",
    required=False,
    default=None,
    help="Run using one of the special modes defined in the configuration file.",
)

parser.add_argument(
    "-f",
    "--fast",
    required=False,
    default=False,
    action="store_true",
    help="Run in fast mode. This only creates the catalogue plots, but does not run any additional scripts.",
)

parser.add_argument(
    "-F",
    "--no-plots",
    required=False,
    default=False,
    action="store_true",
    help="Do not create any plots or web pages. Only produces the autoplotter .yml file. Requires --fast.",
)

if __name__ == "__main__":
    # Parse our lovely arguments and pass them to the velociraptor library
    from velociraptor.autoplotter.objects import AutoPlotter
    from velociraptor.autoplotter.metadata import AutoPlotterMetadata
    from velociraptor.autoplotter.compare import (
        recreate_instances,
        recreate_single_figure,
    )
    from velociraptor import load

    from swiftsimio import load as load_snapshot

    from matplotlib import __version__
    from matplotlib.pyplot import style

    from subprocess import run
    from glob import glob
    import os

    from swiftpipeline.config import Config
    from swiftpipeline.html import WebpageCreator

    PLOT_FILE_EXTENSION = "png"

    args = parser.parse_args()

    if args.no_plots and not args.fast:
        raise AttributeError("Cannot have --no-plots without --fast!")

    is_comparison = len(args.snapshots) > 1
    if args.no_plots and is_comparison:
        raise AttributeError("Cannot have --no-plots when running in comparison mode.")

    # Set up some basic debugging things
    if args.debug:
        from tqdm import tqdm

        # Create a function to print out script runtime statistics as a table
        def print_script_runtime_statistics(script_runtimes):

            # Init variable to compute the sum of scripts' runtimes
            total_runtime = 0.0  # sec

            # Add runtimes of all scripts to the sum
            for script_name, script_runtime in script_runtimes.items():
                if script_name != "Wallclock time":
                    total_runtime += script_runtime

            # Create a key-value pair for the total runtime of the scripts
            script_runtimes["Total CPU time"] = total_runtime

            # Compute the width of first column in the table (add 3 for padding)
            col_width = max(len(name) for name in script_runtimes.keys()) + 3

            # Set the width of the columns
            first_item_width, second_item_width, third_item_width = col_width, 20, 27

            # Make zeroth row
            first_item = "Script name".ljust(first_item_width)
            second_item = "Run time (s)".ljust(second_item_width)
            third_item = "Fraction of total CPU time".ljust(third_item_width)

            # Print out table header including the zeroth row
            print()
            print("Script Run Time Statistics: \n")
            print(f"{first_item}{second_item}{third_item}")
            print("-" * (first_item_width + second_item_width + third_item_width))

            # Create and print out the remaining rows
            for script_name, script_runtime in script_runtimes.items():

                # Fraction of the total runtime per script
                run_time_fraction = script_runtime / script_runtimes["Total CPU time"]

                first_item = f"{script_name}".ljust(first_item_width)
                second_item = f"{script_runtime:.4f}".ljust(second_item_width)
                third_item = f"{run_time_fraction:.6f}".ljust(third_item_width)

                # Print row
                print(f"{first_item}{second_item}{third_item}")

            print()

            return

    def print_if_debug(string: str):
        if args.debug:
            print(string)

    print_if_debug("Running in debug mode. Arguments given are:")
    for name, value in dict(vars(args)).items():
        print_if_debug(f"{name}: {value}")

    config = Config(config_directory=args.config)

    special_mode = None
    if args.special is not None:
        special_mode = config.get_special_mode(args.special)

    print_if_debug(f"Matplotlib version: {__version__}.")
    if config.matplotlib_stylesheet != "default":
        stylesheet_path = f"{config.config_directory}/{config.matplotlib_stylesheet}"
        print_if_debug(f"Applying matplotlib stylesheet at {stylesheet_path}.")
        style.use(stylesheet_path)

    # Reverse so most recently modified is at the top.
    auto_plotter_configs = list(config.auto_plotter_configs)

    box_size_correction_directory = f"{config.config_directory}/box_size_corrections"

    print_if_debug("Loading snapshot metadata")
    snapshots = [
        load_snapshot(f"{input}/{snapshot}")
        for input, snapshot in zip(args.input, args.snapshots)
    ]
    if args.run_names is not None:
        run_names = args.run_names
        print_if_debug("Using custom run names:")
        print_if_debug(" ".join(run_names))
    else:
        # First, check if the snapshots are all at the same redshift
        redshifts = {data.metadata.redshift for data in snapshots}
        # If the size of the set is one, then all redshifts are the same
        if len(redshifts) == 1:
            # All redshifts are the same! No need to modify runs' names
            run_names = [data.metadata.run_name for data in snapshots]
        # If the size of the set > 1, then at least two runs have different redshifts
        else:
            # Need to append appropriate redshifts to names.
            run_names = [
                f"{data.metadata.run_name} (z={data.metadata.redshift:1.3f})"
                for data in snapshots
            ]
        print_if_debug("Using default run names from snapshot:")
        print_if_debug(" ".join(run_names))

    observational_data_path = (
        f"{config.config_directory}/{config.observational_data_directory}/data"
    )

    if not is_comparison:
        # Run the pipeline based on the arguments if only a single simulation is
        # included and generate the metadata yaml file.
        print_if_debug(
            f"Generating initial AutoPlotter instance for {auto_plotter_configs}."
        )

        auto_plotter = AutoPlotter(
            auto_plotter_configs,
            observational_data_directory=observational_data_path,
            correction_directory=box_size_correction_directory,
        )

        halo_catalogue_filename = f"{args.input[0]}/{args.catalogues[0]}"
        print_if_debug(f"Loading halo catalogue at {halo_catalogue_filename}.")

        registration_filename = [
            f"{config.config_directory}/{file_path}"
            for file_path in config.auto_plotter_registration
        ]

        if registration_filename is not None:
            print_if_debug(
                f"Using registration functions contained in {registration_filename}"
            )

        # string pointing to the boolean array within the VelociraptorCatalogue object
        # used to mask the data for all autoplotter plots
        global_mask_tag = config.auto_plotter_global_mask

        if global_mask_tag is not None:
            print_if_debug(f"Masking all catalogue properties using {global_mask_tag}")

        catalogue = load(
            halo_catalogue_filename,
            disregard_units=True,
            registration_file_path=registration_filename,
        )
        if special_mode is not None:
            special_mode.adapt_catalogue(catalogue)
        print_if_debug(f"Linking catalogue and AutoPlotter instance.")
        auto_plotter.link_catalogue(
            catalogue=catalogue, global_mask_tag=global_mask_tag
        )

        print_if_debug(f"Creating figures with extension .png in {args.output}.")
        print_if_debug("Converting AutoPlotter.plots to a tqdm instance.")

        if args.debug:
            auto_plotter.plots = tqdm(auto_plotter.plots, desc="Creating figures")

        auto_plotter.create_plots(
            directory=args.output,
            file_extension=PLOT_FILE_EXTENSION,
            debug=args.debug,
            no_plots=args.no_plots,
        )

        print_if_debug("Creating AutoPlotterMetadata instance.")
        auto_plotter_metadata = AutoPlotterMetadata(auto_plotter=auto_plotter)
        metadata_filename = (
            f"{args.input[0]}/{args.metadata}_{args.snapshots[0][-9:-5]}"
        )
        if special_mode is not None:
            metadata_filename += f"_{special_mode.name}"
        metadata_filename += ".yml"
        print_if_debug(f"Creating and writing metadata to {metadata_filename}")

        try:
            auto_plotter_metadata.write_metadata(metadata_filename)
        except (OSError, PermissionError) as e:
            print_if_debug(f"Unable to save metadata to {metadata_filename}")
            pass
    else:
        # Need to generate our data again from the built-in yaml files.
        if special_mode is not None:
            suffix = f"_{special_mode.name}.yml"
        else:
            suffix = ".yml"
        metadata_filenames = [
            f"{input}/{args.metadata}_{snapshot[-9:-5]}{suffix}"
            for input, snapshot in zip(args.input, args.snapshots)
        ]

        if args.debug:
            for metadata_filename in metadata_filenames:
                if not os.path.exists(metadata_filename):
                    print_if_debug(
                        f"Unable to find {metadata_filename}, ensure the pipeline has "
                        "been run in standalone mode for this simulation before "
                        "attempting comparisons."
                    )

        print_if_debug(f"Attempting to recreate instances for {metadata_filenames}")
        auto_plotter, auto_plotter_metadata, line_data = recreate_instances(
            config=auto_plotter_configs,
            paths=metadata_filenames,
            names=run_names,
            observational_data_directory=observational_data_path,
            file_extension=PLOT_FILE_EXTENSION,
            correction_directory=box_size_correction_directory,
        )

        if not os.path.exists(args.output):
            os.mkdir(args.output)

        print_if_debug("Converting AutoPlotter.plots to a tqdm instance.")
        if args.debug:
            auto_plotter.plots = tqdm(auto_plotter.plots, desc="Creating figures")

        for plot in auto_plotter.plots:
            try:
                recreate_single_figure(
                    plot=plot,
                    line_data=line_data,
                    output_directory=args.output,
                    file_type=PLOT_FILE_EXTENSION,
                )
            except:
                print_if_debug(f"Unable to create figure {plot.filename}.")

    # Now that we have auto_plotter_metadata we can use it to check if we have
    # inadvertently created multiple plots with the same filename.
    if args.debug:
        figure_filenames = {plot.filename: 0 for plot in auto_plotter_metadata.plots}

        for plot in auto_plotter_metadata.plots:
            figure_filenames[plot.filename] += 1

        for filename, number_of_figures in figure_filenames.items():
            if number_of_figures > 1:
                print_if_debug(
                    f"{number_of_figures} figures with filename "
                    f"{filename}.{PLOT_FILE_EXTENSION} have been created. "
                    "This will cause overwriting, and may not be intentional."
                )

    # Now move onto using the ``config`` to generate plots from the actual data.
    scripts_to_use = config.comparison_scripts if is_comparison else config.scripts

    full_script_path_list = [
        f"{config.config_directory}/{script.filename}" for script in scripts_to_use
    ]
    script_additional_args_list = [
        script.additional_argument_list for script in scripts_to_use
    ]

    def script_run(script_path: str, script_args: List[Any]) -> Tuple[str, float]:
        """
        A function through which scripts can be run in parallel using the p_map method
        from p_tqdm library.

        Parameters
        ----------

        script_path: str
            Absolute path to the script executable including the file name

        script_args: List[Any]
            Optional arguments to the script

        Returns
        -------

        output: Tuple[str, float]
            A tuple containing the script (file) name and the time it took to
            run the script.

        """

        time_start = time()
        run(
            [
                "python3",
                script_path,
                "-s",
                *args.snapshots,
                "-c",
                *args.catalogues,
                "-d",
                *args.input,
                "-n",
                *run_names,
                "-o",
                args.output,
                "-C",
                config.config_directory,
                *script_args,
            ]
        )
        time_end = time()

        # Record the time difference before and after the call of subprocess.run
        script_runtime = time_end - time_start

        try:
            script_name = script_path.split("/")[-1]
        except (AttributeError, TypeError):
            script_name = "Other calls"

        return (script_name, script_runtime)

    if args.fast:
        print("Fast mode, not running scripts.")
    else:
        # Time the call of p_map to measure total wall clock time
        time_start = time()

        if args.num_of_cpus is not None:

            print_if_debug(
                f"Running scripts in parallel using {args.num_of_cpus} CPUs."
            )

            script_runs_out = p_map(
                script_run,
                full_script_path_list,
                script_additional_args_list,
                num_cpus=args.num_of_cpus,
                desc="Running Scripts",
            )
        else:

            print_if_debug(
                f"Running scripts in parallel on the maximum number of CPUs "
                f"avaliable."
            )

            script_runs_out = p_map(
                script_run,
                full_script_path_list,
                script_additional_args_list,
                desc="Running Scripts",
            )

        # Finish calculating wall clock time
        time_end = time()

        # Print out script runtime statistics
        if args.debug:

            # Record wall clock time used to execute all scripts
            wall_clock_time = time_end - time_start

            # Create dict for recording CPU time spent while running scripts
            script_runtimes = {}

            # Convert list of tuples to dict
            for (name, time) in script_runs_out:
                script_runtimes[name] = time

            # Add wallclock time
            script_runtimes["Wallclock time"] = wall_clock_time

            # Print out runtime scatistics
            print_script_runtime_statistics(script_runtimes=script_runtimes)

    if not args.no_plots:
        # Create the webpage`
        print_if_debug("Creating webpage.")
        warnings = []
        if args.fast:
            warnings.append(
                f"Running in fast mode: only autoplotter plots are included."
            )
        if special_mode is not None:
            warnings.append(f"Running in special mode {special_mode.name}.")
        webpage = WebpageCreator(warnings)
        webpage.add_auto_plotter_metadata(auto_plotter_metadata=auto_plotter_metadata)
        if not args.fast:
            webpage.add_config_metadata(config=config, is_comparison=is_comparison)
        page_name = " | ".join(run_names)
        if args.fast:
            page_name = f"FAST mode - {page_name}"
        if special_mode is not None:
            page_name = f"SPECIAL {special_mode.name} - {page_name}"
        webpage.add_metadata(page_name=page_name)
        webpage.add_run_metadata(config=config, snapshots=snapshots)
        webpage.render_webpage()
        webpage.save_html(f"{args.output}/index.html")

    print_if_debug("Done.")