plot_ensemble_parameter_single.py

"""
Analyse a PCSE ensemble with one varying parameter, as generated by wofost_ensemble_parameters.py.

Example:
    %run plot_ensemble_parameter_single.py outputs/RDMSOL -v
"""
from tqdm import tqdm

import fpcup

### Parse command line arguments
import argparse
parser = argparse.ArgumentParser(description="Analyse a PCSE ensemble with one varying parameter, as generated by wofost_ensemble_parameters.py.")
parser.add_argument("output_dir", help="folder to load PCSE outputs from", type=fpcup.io.Path)
parser.add_argument("--results_dir", help="folder to save plots into", type=fpcup.io.Path, default=fpcup.DEFAULT_RESULTS/"sensitivity")
parser.add_argument("-v", "--verbose", help="increase output verbosity", action="store_true")
args = parser.parse_args()

args.PARAMETER_NAME = args.output_dir.stem

### This gets executed only when the script is run normally; not by multiprocessing.
if __name__ == "__main__":
    fpcup.multiprocessing.freeze_support()

    ### SETUP
    # Load the ensemble summary
    summary = fpcup.io.load_combined_ensemble_summary_geo(args.output_dir)
    summary.sort_values(args.PARAMETER_NAME, inplace=True)
    if args.verbose:
        print(f"Loaded summary from {args.output_dir.absolute()}")


    ### PLOTTING
    # Loop over the crops and generate a figure for each
    for (crop_name, summary_by_crop) in tqdm(summary.groupby("crop"), desc="Plotting figures", unit="crop", leave=args.verbose):
        # Setup filename
        crop_short = fpcup.crop.CROP2ABBREVIATION[crop_name]
        saveto = args.results_dir/f"{args.PARAMETER_NAME}-{crop_short}.pdf"

        # Plot
        fpcup.plotting.sensitivity_one_crop(summary_by_crop, crop_name, args.PARAMETER_NAME, saveto=saveto)