pipeline/pipeline_disk.py

# coding: utf-8

import os
from shutil import copy
import argparse
import time
import numpy as np
import disk_postprocess as dp
from functools import reduce

storage_in = "/home/nbrucy/simus/"
storage_out = "/home/nbrucy/visus/"

parser = argparse.ArgumentParser()

parser.add_argument("runs", help="name of runs", nargs="*", default=["015_iso"])
parser.add_argument("-b", "--begin", help="id of first output", type=int, default=1)
parser.add_argument("-e", "--end", help="id of last output", type=int, default=100)
parser.add_argument("-s", "--step", help="step between two output", type=int, default=1)
parser.add_argument("-p", "--project", help="specify project name", default="disk")
parser.add_argument(
    "-fr",
    "--force_redo",
    help="redo plots even if the files already exist",
    action="store_true",
)
parser.add_argument(
    "-w", "--watch", help="wait and watch for missing outputs", action="store_true"
)
parser.add_argument("--skip", help="skip failed loadings", action="store_true")
parser.add_argument(
    "-wt",
    "--waiting_time",
    help="time between to successive try when watching new outputs (in second)",
    type=int,
    default=120,
)
parser.add_argument(
    "-af",
    "--allowed_failures",
    help="number of allowed failures when waiting",
    default=30,
)
parser.add_argument("-i", "--interactive", help="Interactive mode", action="store_true")


parser.add_argument(
    "-d", "--disk", help="compute specific disk radial analysis", action="store_true"
)
parser.add_argument(
    "-pd", "--plot_disk", help="plot specific disk radial analysis", action="store_true"
)
parser.add_argument("-m", "--maps", help="compute generic maps", action="store_true")
parser.add_argument("-pm", "--plot_maps", help="plot generic maps", action="store_true")
parser.add_argument(
    "-c", "--compare", help="compare different runs", action="store_true"
)
parser.add_argument(
    "-ev",
    "--evolution",
    help="plot evolution of quantities over time",
    action="store_true",
)
parser.add_argument(
    "--pdf", help="plot pdf of fluctuations of column density", action="store_true"
)

parser.add_argument(
    "--fft", help="use quick and dirty fft rendering", action="store_true"
)
parser.add_argument(
    "--images",
    nargs="*",
    default=["coldens", "rho", "speed", "Q", "T"],
    choices=["coldens", "rho", "speed", "Q", "T", "levels", "cpu"],
)
parser.add_argument(
    "--axes", nargs="*", default=["x", "y", "z"], choices=["x", "y", "z"]
)
parser.add_argument("--zoom", help="zoom", type=float, default=2.0)
parser.add_argument(
    "-ms",
    "--mapsize",
    help="size of the maps created in he map mode (in pixel)",
    type=int,
    default=1024,
)


parser.add_argument(
    "--nb_bin", help="Number of bins for azimuthal averages", type=int, default=50
)
parser.add_argument(
    "--binning",
    help="Kind of binning (logarithmic or linear)",
    choices=["log", "lin"],
    default="log",
)
parser.add_argument(
    "--rad_ext", help="Value of the highest bin", type=float, default=1.0
)
parser.add_argument(
    "-x", help="x position of the central point", type=float, default=1.0
)
parser.add_argument(
    "-y", help="y position of the central point", type=float, default=1.0
)
parser.add_argument(
    "-z", help="z position of the central point", type=float, default=1.0
)

parser.add_argument(
    "-ta", "--time_avg", help="Plot time averaged comparaison", action="store_true"
)


parser.add_argument(
    "--colormap", help="Colormap used", choices=dp.P.colormaps(), default="plasma"
)
parser.add_argument(
    "--format",
    help="Format of the output",
    choices=["png", "jpeg", "pdf", "svg", "ps"],
    default="jpeg",
)
parser.add_argument("--dpi", help="Resolution of the output", type=int, default=400)
parser.add_argument("--beamer", help="Beamer mode", action="store_true")


args = parser.parse_args()

project = args.project
runs = args.runs

first = args.begin
last = args.end
step = args.step

rad = 0.5 / args.zoom

# extension for out files
dp.out_ext = "." + args.format
if format == "pdf":
    dp.P.style.use("pdf")

if args.beamer:
    dp.P.rcParams["font.family"] = "sans-serif"
    dp.P.rcParams["figure.figsize"] = (5, 3.5)

# Plot properties
dp.P.rcParams["image.cmap"] = args.colormap
dp.P.rcParams["savefig.dpi"] = args.dpi

# List of id that were successfully computed
run_succeded = {}

for run in runs:
    path_suffix = project + "/" + run
    path_in = storage_in + path_suffix
    path_out = storage_out + path_suffix

    if not os.path.exists(path_out):
        os.makedirs(path_out)
        try:
            copy(path_in + "/disk.nml", path_out)
            copy(path_in + "/output_00001/compilation.txt", path_out)
        except:
            pass

    run_succeded[run] = []

    for i in range(first, last + 1, step):
        failures = 0
        success = False

        while not success:
            try:
                maps_disk = None
                if args.maps:
                    print("[{}, {}] computing maps".format(run, i))
                    maps_disk = dp.compute_image_data(
                        path_in,
                        i,
                        radius=rad,
                        path_out=path_out,
                        tag=run,
                        map_size=args.mapsize,
                        force=args.force_redo,
                        axes_los=args.axes,
                        images=args.images,
                        pos_star=np.array([args.x, args.y, args.z]),
                        fft=args.fft,
                    )
                    print("[{}, {}] maps computed".format(run, i))
                if args.plot_maps:
                    print("[{}, {}] plotting maps".format(run, i))
                    maps_disk = dp.plot_maps(
                        path_out,
                        i,
                        maps_disk=maps_disk,
                        axes_los=args.axes,
                        images=args.images,
                        tag=run,
                        force=args.force_redo,
                        interactive=args.interactive,
                        put_title=(not args.beamer),
                    )
                    print("[{}, {}] maps plotted".format(run, i))
                if args.disk:
                    print("[{}, {}] computing disk props".format(run, i))
                    dp.disk_prop(
                        path_in,
                        i,
                        path_out=path_out,
                        nb_bin=args.nb_bin,
                        binning=args.binning,
                        rad_ext=args.rad_ext,
                        force=args.force_redo,
                        pos_star=np.array([args.x, args.y, args.z]),
                    )
                    print("[{}, {}] disk_props computed".format(run, i))
                if args.plot_disk:
                    print("[{}, {}] plotting disk props".format(run, i))
                    dp.plot_disk_prop(
                        path_out,
                        i,
                        tag=run,
                        force=args.force_redo,
                        interactive=args.interactive,
                        put_title=(not args.beamer),
                    )
                    print("[{}, {}] disk props plotted".format(run, i))
                if args.pdf:
                    print("[{}, {}] computing pdf".format(run, i))
                    dp.disk_pdf(
                        path_out,
                        i,
                        maps_disk,
                        pos_star=np.array([args.x, args.y, args.z]),
                        force=args.force_redo,
                        tag=run,
                        interactive=args.interactive,
                        put_title=(not args.beamer),
                    )
                    print("[{}, {}] pdf computed".format(run, i))
                # If we are here, success !
                success = True
                run_succeded[run].append(i)
            except (ValueError, IOError, KeyError) as e:
                print(e)
                if args.watch and failures < args.allowed_failures:
                    failures = failures + 1
                    print(
                        "Unable to proceed for run {} output {}. Trying again in {} s ({} tries remaining)".format(
                            run, i, args.waiting_time, args.allowed_failures - failures
                        )
                    )
                    time.sleep(args.waiting_time)
                elif args.skip:
                    break
                else:
                    raise
    if args.evolution:
        print("[{}] plotting evolution".format(run))
        dp.evolution(
            path_out,
            run_succeded[run],
            force=args.force_redo,
            interactive=args.interactive,
            pdf=args.pdf,
        )
        print("[{}] evolution plotted".format(run))

if args.compare:
    path_suffix = project
    path = storage_out + path_suffix

    if args.time_avg:
        # Select output availables for at least on run
        output_ok = reduce(np.union1d, [run_succeded[run] for run in runs]).astype(int)
        print(output_ok)
        if len(output_ok) >= 1:
            dp.compare(
                path,
                runs,
                output_ok,
                path_out=path + "/comp",
                force=args.force_redo,
                interactive=args.interactive,
                Q_in_name=(not args.pdf),
                pdf=args.pdf,
            )
    else:
        for i in range(first, last + 1, step):
            # Select usable runs
            runs_ok = [run for run in runs if i in run_succeded[run]]

            try:
                if len(runs_ok) > 1:
                    dp.compare(
                        path,
                        runs_ok,
                        i,
                        path_out=path + "/comp",
                        force=args.force_redo,
                        interactive=args.interactive,
                        Q_in_name=(not args.pdf),
                        pdf=args.pdf,
                    )
            except (KeyError, IOError) as e:
                print(e)
                if args.skip:
                    pass
                else:
                    raise