diff --git a/workflows/evolutive_controller/postprocess.sh b/workflows/evolutive_controller/postprocess.sh new file mode 100755 index 00000000..6d1ec030 --- /dev/null +++ b/workflows/evolutive_controller/postprocess.sh @@ -0,0 +1,17 @@ +#!/bin/bash +# Validation plots: reference vs NICE output for R, Z, Ip, PF coil current. Run by +# bin/pds-run-case.sbatch after muscle_manager finishes -- see that script's header for the +# PDS_REPO/SCENARIOS_REPO/SHOT/CASE_DIR/RUN_DIR/PYTHON contract. +# +# reconstruction_uri mirrors settings.ymmsl's source.source_uri (the inverse_convergence run +# this case's F_INIT reference equilibrium is read from). +set -euo pipefail + +mkdir -p "$RUN_DIR/plots" + +"$PYTHON" "$PDS_REPO/workflows/utils/plot_validation_evolutive_controller.py" \ + --shot_nr "$SHOT" \ + --dina_uri "$SCENARIOS_REPO/$SHOT/data/in" \ + --reconstruction_uri "$PDS_REPO/cases/runs/inverse_convergence_${SHOT}/out_nice" \ + --nice_uri "$RUN_DIR/out_nice" \ + --output_dir "$RUN_DIR/plots" diff --git a/workflows/utils/plot_validation_evolutive_controller.py b/workflows/utils/plot_validation_evolutive_controller.py new file mode 100644 index 00000000..05f54966 --- /dev/null +++ b/workflows/utils/plot_validation_evolutive_controller.py @@ -0,0 +1,218 @@ +""" +Validate evolutive_controller output: reference vs NICE for R, Z, Ip, PF coil current. +""" + +import argparse +import logging + +import matplotlib.pyplot as plt +import numpy as np +from imas import DBEntry +from imas.ids_defs import CLOSEST_INTERP + +logger = logging.getLogger(__name__) + +PLOT_KWARGS = {"marker": "."} +GET_KWARGS = {"interpolation_method": CLOSEST_INTERP, "lazy": True} + + +def handle_args(): + parser = argparse.ArgumentParser( + description="Compare evolutive_controller's NICE output against its reference" + ) + parser.add_argument("--shot_nr", type=str, help="Shot number") + parser.add_argument( + "--dina_uri", + type=str, + help="URI to load raw DINA scenario data from (reference Ip and PF coil current)", + ) + parser.add_argument( + "--reconstruction_uri", + type=str, + help=( + "URI to load the NICE-reconstructed equilibrium the workflow's `source` actor " + "reads (reference R, Z -- waveform_editor never overlays geometric_axis, so " + "this is the actual reference the controller solves against)" + ), + ) + parser.add_argument( + "--nice_uri", + type=str, + help="URI to load evolutive_controller's NICE output from", + ) + parser.add_argument( + "--output_dir", type=str, help="path to directory in which to put plots" + ) + args = parser.parse_args() + return args + + +def interp_to_range(t_min, t_max, t, *values): + """Restrict a time series to [t_min, t_max], interpolating value(s) exactly at the + boundaries so a coarsely-sampled series draws a line spanning the same x-range as the + series it's compared against, instead of overshooting to its own native samples + outside that range.""" + inside = (t > t_min) & (t < t_max) + t_new = np.concatenate(([t_min], t[inside], [t_max])) + values_new = tuple( + np.concatenate(([np.interp(t_min, t, v)], v[inside], [np.interp(t_max, t, v)])) + for v in values + ) + return (t_new, *values_new) + + +def align_ranges(t_a, values_a, t_b, values_b): + """Restrict both series to whichever of the two spans less time, interpolating the + other series exactly at that shorter span's boundaries so both lines drawn cover the + same x-range.""" + if (t_a.max() - t_a.min()) <= (t_b.max() - t_b.min()): + t_min, t_max = t_a.min(), t_a.max() + else: + t_min, t_max = t_b.min(), t_b.max() + return ( + interp_to_range(t_min, t_max, t_a, *values_a), + interp_to_range(t_min, t_max, t_b, *values_b), + ) + + +def nice_output_flags(db): + """Per-slice NICE solver status: -1 means NICE failed to converge that slice.""" + equilibrium = db.get("equilibrium", lazy=True) + flags = equilibrium.code.output_flag + if not flags: + return np.zeros(len(equilibrium.time)) + return np.asarray(flags) + + +def main(): + """Plot evolutive_controller reference-vs-NICE validation figures""" + args = handle_args() + dbs = { + "dina": DBEntry(f"imas:hdf5?path={args.dina_uri}", "r"), + "reconstruction": DBEntry(f"imas:hdf5?path={args.reconstruction_uri}", "r"), + "nice": DBEntry(f"imas:hdf5?path={args.nice_uri}", "r"), + } + + rz_ip_plot(args, dbs) + kcurr_plot(args, dbs) + + for db in dbs.values(): + db.close() + + +def rz_ip_plot(args, dbs): + """Plot R, Z (plasma boundary geometric axis) and Ip: reference vs NICE output. + + R/Z reference comes from the reconstructed equilibrium fed into `source` (waveform_editor + only overlays boundary/outline and global_quantities/ip from DINA, never geometric_axis -- + see workflows/evolutive_controller/waveforms.yaml), Ip reference from raw DINA. + """ + figure_path = f"{args.output_dir}/pds_rz_ip_{args.shot_nr}.png" + + recon_eq = dbs["reconstruction"].get("equilibrium") + dina_eq = dbs["dina"].get("equilibrium") + nice_eq = dbs["nice"].get("equilibrium") + + recon_t = np.asarray(recon_eq.time) + recon_r = np.asarray([ts.boundary.geometric_axis.r for ts in recon_eq.time_slice]) + recon_z = np.asarray([ts.boundary.geometric_axis.z for ts in recon_eq.time_slice]) + + dina_t = np.asarray(dina_eq.time) + dina_ip = np.asarray([ts.global_quantities.ip for ts in dina_eq.time_slice]) + + nice_mask = nice_output_flags(dbs["nice"]) != -1 + nice_t = np.asarray(nice_eq.time)[nice_mask] + nice_r = np.asarray([ts.boundary.geometric_axis.r for ts in nice_eq.time_slice])[ + nice_mask + ] + nice_z = np.asarray([ts.boundary.geometric_axis.z for ts in nice_eq.time_slice])[ + nice_mask + ] + nice_ip = np.asarray([ts.global_quantities.ip for ts in nice_eq.time_slice])[ + nice_mask + ] + + (recon_t, recon_r, recon_z), (nice_rz_t, nice_r, nice_z) = align_ranges( + recon_t, (recon_r, recon_z), nice_t, (nice_r, nice_z) + ) + (dina_t, dina_ip), (nice_ip_t, nice_ip) = align_ranges( + dina_t, (dina_ip,), nice_t, (nice_ip,) + ) + + fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(18, 5)) + fig.suptitle(f"{args.shot_nr}: reference vs NICE output", fontsize=16) + + panels = [ + (axes[0], "R [m]", recon_t, recon_r, nice_rz_t, nice_r), + (axes[1], "Z [m]", recon_t, recon_z, nice_rz_t, nice_z), + (axes[2], "Ip [A]", dina_t, dina_ip, nice_ip_t, nice_ip), + ] + for ax, label, ref_t, ref_v, out_t, out_v in panels: + ax.set_title(label) + ax.set_ylabel(label) + ax.set_xlabel("time") + ax.plot(ref_t, ref_v, label="reference", **PLOT_KWARGS) + ax.plot(out_t, out_v, label="nice", **PLOT_KWARGS) + ax.legend() + + fig.tight_layout(rect=(0, 0.03, 1, 0.95)) + fig.savefig(figure_path) + + +def kcurr_plot(args, dbs): + """Plot PF coil current (Kcurr, the KCURR_PFPO1 controller's tracked quantity): + the raw-DINA target vs NICE's actual per-coil current, one panel per coil.""" + coil_figure_path = f"{args.output_dir}/pds_kcurr_{args.shot_nr}.png" + coil_dict = {} + pfas = { + "reference": dbs["dina"].get("pf_active"), + "nice": dbs["nice"].get("pf_active"), + } + nice_mask = nice_output_flags(dbs["nice"]) != -1 + ref_time = np.asarray(pfas["reference"].time) + nice_time = np.asarray(pfas["nice"].time)[nice_mask] + if ref_time.max() - ref_time.min() <= nice_time.max() - nice_time.min(): + t_min, t_max = ref_time.min(), ref_time.max() + else: + t_min, t_max = nice_time.min(), nice_time.max() + + nrows, ncols = (7, 2) + fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(15, 22)) + fig.suptitle(f"{args.shot_nr}: reference vs NICE PF coil current", fontsize=16) + axes = axes.flatten() + + for key, pfa in pfas.items(): + mask = nice_mask if key == "nice" else None + full_time = np.asarray(pfa.time) + for coil in pfa.coil: + coil_name = str(coil.name) + if coil_name not in coil_dict: + next_slot = max(coil_dict.values(), default=-1) + 1 + if next_slot >= len(axes): + logger.warning( + "pf_active coil name %r (%s) has no free plot slot " + "(reference/nice disagree on coil naming for this " + "scenario) -- skipping its plot.", + coil_name, + key, + ) + continue + coil_dict[coil_name] = next_slot + axes[coil_dict[coil_name]].set_title(coil_name) + axes[coil_dict[coil_name]].set_ylabel("current") + axes[coil_dict[coil_name]].set_xlabel("time") + time, current = full_time, np.asarray(coil.current.data) + if mask is not None and len(mask) == len(time): + time, current = time[mask], current[mask] + time, current = interp_to_range(t_min, t_max, time, current) + axes[coil_dict[coil_name]].plot(time, current, label=key, **PLOT_KWARGS) + axes[coil_dict[coil_name]].legend() + for ax in axes[len(pfa.coil) :]: + fig.delaxes(ax) + + fig.tight_layout(rect=(0, 0.03, 1, 0.95)) + fig.savefig(coil_figure_path) + + +if __name__ == "__main__": + main()