swift-image

#!python3
"""
SWIFT-image is a companion to SWIFT-pipeline that allows the user to
create an array of images of a given simulation.
"""

import argparse as ap
from typing import Union

parser = ap.ArgumentParser(
    prog="swift-image",
    description=(
        "Creates a webpage containing many images, based on the metadata files "
        "provided alongside swift-pipeline. Note that, unlike swift-pipeline, "
        "swift-image cannot be used for comparisons. Do not create the images in "
        "the same directory as the pipeline outputs, as this will overwrite your "
        "index.html."
    ),
    epilog=(
        "Example usage:\n"
        "swift-pipeline -C ~/config -c example_0000.properties -s snapshot_0000.hdf5 "
        "-o ~/plots/example_0000/images -i /path/to/my/sim\n\n"
    ),
)

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

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

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

parser.add_argument(
    "-o",
    "--output",
    type=str,
    required=True,
    help="Output directory for images. 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."
    ),
)

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(
    "-p",
    "--parallel",
    required=False,
    default=False,
    action="store_true",
    help="Run in parallel if flag is present. May not work on all systems. Default: no.",
)

if __name__ == "__main__":
    # Parse our lovely arguments and pass them to the velociraptor library
    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.imageconfig import ImageConfig
    from swiftpipeline.imaging import create_all_images

    from pathlib import Path

    args = parser.parse_args()

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

    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)
    image_config = ImageConfig(config_directory=args.config)

    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)

    output_path = Path(args.output)
    output_path.mkdir(exist_ok=True, parents=True)

    input_path = Path(args.input)
    snapshot_path = input_path / args.snapshot
    catalogue_path = input_path / args.catalogue

    create_all_images(
        config=image_config,
        output_path=output_path,
        snapshot_path=snapshot_path,
        catalogue_path=catalogue_path,
        parallel=args.parallel,
        debug=args.debug,
    )

    print_if_debug("Done.")