Refactor sensitivity plot

demo.ipynb

@@ -662,7 +662,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "n = 51  # Number of values to test\n",
+    "n = 101  # Number of values to test\n",
     "variable_names = [\"WAV\"]\n",
     "combined_parameters = fpcup.parameters.generate_ensemble_space(*variable_names, n=n)\n",
     "variables = {\"override\": combined_parameters}\n",
@@ -697,7 +697,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# TO DO: Refactor ensemble plot"
+    "summary = fpcup.io.load_combined_ensemble_summary_geo(output_dir)\n",
+    "fpcup.plotting.sensitivity_one_crop(summary, crop_name, variable_names[0])"
    ]
   },
   {

fpcup/plotting/__init__.py

@@ -3,4 +3,5 @@
 
 from .brp import *
 from .brp_folium import *
+from .sensitivity import *
 from .wofost import *

fpcup/plotting/sensitivity.py

@@ -0,0 +1,79 @@
+"""
+Functions for plotting data/results relating to sensitivity studies.
+"""
+import pandas as pd
+
+from matplotlib import pyplot as plt
+
+from ..model import GeoSummary
+from ..parameters import PCSEParameter, all_parameters, pcse_summary_outputs
+from ..typing import Optional, PathOrStr
+
+### CONSTANTS
+OUTPUT_LABELS = ("DOM", "RD", "LAIMAX", "TAGP", "TWSO")  # Parameters to plot
+OUTPUT_PARAMETERS = [pcse_summary_outputs[key] for key in OUTPUT_LABELS]
+
+
+### HELPER FUNCTIONS
+def number_of_replicates(data: pd.DataFrame, column: str) -> int:
+    """
+    Determine the number of replicates in a given column, e.g. number_of_replicates(data, column='geometry') is the number of sites.
+    """
+    return len(data[column].unique())
+
+
+def sensitivity_one_crop(summary_by_crop: GeoSummary, crop_name: str, parameter: str | PCSEParameter, *,
+                         saveto: Optional[PathOrStr]=None) -> None:
+    """
+    For a summary table for one crop, plot the sensitivity of certain output parameters to the given input parameter.
+    """
+    if not isinstance(parameter, PCSEParameter):
+        parameter = all_parameters[parameter]
+
+    # Determine number of runs
+    n_sites = number_of_replicates(summary_by_crop, "geometry")
+    n_soiltypes = number_of_replicates(summary_by_crop, "soiltype")
+    n_parameter_values = number_of_replicates(summary_by_crop, parameter.name)
+
+    # Setup
+    fig, axs = plt.subplots(nrows=len(OUTPUT_PARAMETERS), ncols=n_soiltypes, sharex=True, sharey="row", figsize=(3*n_soiltypes, 10), squeeze=False)
+
+    # Loop over the columns (soil types) first
+    for ax_col, (soiltype, summary_by_soiltype) in zip(axs.T, summary_by_crop.groupby("soiltype")):
+        ax_col[0].set_title(f"Soil type: {soiltype}")
+
+        # Loop over the rows (summary outputs) next
+        for ax, output in zip(ax_col, OUTPUT_PARAMETERS):
+            # Plot a line for each site
+            summary_by_soiltype.groupby("geometry").plot.line(parameter.name, output.name, ax=ax, alpha=0.5, legend=False)
+
+    # Add reference line for default value, if available
+    try:
+        for ax in axs.ravel():
+            ax.axvline(parameter.default, color="black", linestyle="--", alpha=0.5)
+    except (TypeError, AttributeError):
+        pass
+
+    # Titles / labels
+    try:
+        xlim = parameter.bounds
+    except AttributeError:
+        xlim = summary_by_crop[parameter.name].min(), summary_by_crop[parameter.name].max()
+    axs[0, 0].set_xlim(*xlim)
+
+    for ax in axs[1:, 0]:
+        ax.set_ylim(ymin=0)
+
+    for ax, output in zip(axs[:, 0], OUTPUT_PARAMETERS):
+        ax.set_ylabel(output.name)
+
+    fig.suptitle(f"WOFOST sensitivity to {parameter.name}: {crop_name} ({n_sites} sites, {n_parameter_values} values)\n{parameter}")
+    fig.align_xlabels()
+    fig.align_ylabels()
+
+    # Save figure
+    if saveto is not None:
+        plt.savefig(saveto, bbox_inches="tight")
+    else:
+        plt.show()
+    plt.close()

plot_ensemble_parameter_single.py

@@ -4,7 +4,6 @@
 Example:
     %run plot_ensemble_parameter_single.py outputs/RDMSOL -v
 """
-from matplotlib import pyplot as plt
 from tqdm import tqdm
 
 import fpcup
@@ -18,89 +17,25 @@
 args = parser.parse_args()
 
 args.PARAMETER_NAME = args.output_dir.stem
-args.PARAMETER = fpcup.parameters.all_parameters[args.PARAMETER_NAME]
-
-
-### CONSTANTS
-OUTPUT_LABELS = ("DOM", "RD", "LAIMAX", "TAGP", "TWSO")  # Parameters to plot
-OUTPUT_PARAMETERS = [fpcup.parameters.pcse_summary_outputs[key] for key in OUTPUT_LABELS]
-
-
-### HELPER FUNCTIONS
-def number_of_replicates(data: fpcup.geo.gpd.GeoDataFrame, column: str) -> int:
-    """
-    Determine the number of replicates in a given column, e.g. number_of_replicates(data, column='geometry') is the number of sites.
-    """
-    return len(data[column].unique())
-
 
 ### 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
-    inputsummary, summary = fpcup.io.load_ensemble_summary_from_folder(args.output_dir)
-    if args.verbose:
-        print(f"Loaded input summary, summary from {args.output_dir.absolute()}")
-
-    # Join
-    inputsummary.drop(columns=["geometry"], inplace=True)
-    summary = summary.join(inputsummary)
+    summary = fpcup.io.load_combined_ensemble_summary_geo(args.output_dir)
     summary.sort_values(args.PARAMETER_NAME, inplace=True)
     if args.verbose:
-        print("Joined input/output summary tables")
-
-    # Determine number of runs
-    n_crops = number_of_replicates(summary, "crop")
+        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"
 
-        # Determine number of runs
-        n_sites = number_of_replicates(summary_by_crop, "geometry")
-        n_soiltypes = number_of_replicates(summary_by_crop, "soiltype")
-        n_parameter_values = number_of_replicates(summary_by_crop, args.PARAMETER_NAME)
-
-        # Setup
-        fig, axs = plt.subplots(nrows=len(OUTPUT_PARAMETERS), ncols=n_soiltypes, sharex=True, sharey="row", figsize=(10, 10), squeeze=False)
-
-        # Loop over the columns (soil types) first
-        for ax_col, (soiltype, summary_by_soiltype) in zip(axs.T, summary_by_crop.groupby("soiltype")):
-            ax_col[0].set_title(f"Soil type: {soiltype}")
-
-            # Loop over the rows (summary outputs) next
-            for ax, output in zip(ax_col, OUTPUT_PARAMETERS):
-                # Plot a line for each site
-                summary_by_soiltype.groupby("geometry").plot.line(args.PARAMETER_NAME, output.name, ax=ax, alpha=0.5, legend=False)
-
-        # Add reference line for default value, if available
-        try:
-            for ax in axs.ravel():
-                ax.axvline(args.PARAMETER.default, color="black", linestyle="--", alpha=0.5)
-        except (TypeError, AttributeError):
-            pass
-
-        # Titles / labels
-        try:
-            xlim = args.PARAMETER.bounds
-        except AttributeError:
-            xlim = summary_by_crop[args.PARAMETER_NAME].min(), summary_by_crop[args.PARAMETER_NAME].max()
-        axs[0, 0].set_xlim(*xlim)
-
-        for ax in axs[1:, 0]:
-            ax.set_ylim(ymin=0)
-
-        for ax, output in zip(axs[:, 0], OUTPUT_PARAMETERS):
-            ax.set_ylabel(output.name)
-
-        fig.suptitle(f"WOFOST sensitivity to {args.PARAMETER_NAME}: {crop_name} ({n_sites} sites, {n_parameter_values} values)\n{args.PARAMETER}")
-        fig.align_xlabels()
-        fig.align_ylabels()
-
-        # Save figure
-        plt.savefig(args.results_dir/f"{args.PARAMETER_NAME}-{crop_short}.pdf", bbox_inches="tight")
-        plt.close()
+        # Plot
+        fpcup.plotting.sensitivity_one_crop(summary_by_crop, crop_name, args.PARAMETER_NAME, saveto=saveto)