Cleanup TSG analysis notebooks

analysis.py

@@ -1,13 +1,28 @@
 """Helper functions used by notebooks/analysis*.ipynb files."""
 from __future__ import annotations
 
+from collections.abc import Sequence
 from typing import Any
 
 import matplotlib.pyplot as plt
 import numpy as np
 
+from utils import savefig
 
-V_MIN, V_MAX = (11.4**2, 12.6**2)  # +/-5%, units kV^2
+
+V_MIN, V_MAX = 11.4, 12.6  # 12 ± 5%, units kV
+V_MIN2, V_MAX2 = V_MIN**2, V_MAX**2  # units kV^2
+
+# Match Fig8 in Qu and Li (2020)
+# - they count the substation as bus 1
+# - we count the substation as bus 0
+BUSES = (8, 18, 21, 30, 39, 45, 54)  # 0 = substation
+
+TIME_TICKS =  (   0, 2400, 4800,  7200,  9600, 12000, 14400)
+TIME_LABELS = ('0h', '4h', '8h', '12h', '16h', '20h', '24h')
+
+Y_MIN, Y_MAX = 11.2, 12.8  # good for plotting V_MIN to V_MAX
+YTICKS = (11.4, 11.7, 12, 12.3, 12.6)
 
 
 def calculate_violations(
@@ -35,29 +50,29 @@ def calculate_violations(
     - 'max viol': among all bus-timestep violations, the maximum absolute magnitude
         of the violation
 
-    Args:
-        key: tuple (info_provided, seed), where seed may be None when info_provided == 'known'
-        pkl: saved results read from pickle
-        ax: optional axes to plot a histogram of violations
-
-    Returns:
-        num_mistakes: total number of mistakes
-        avg_viol: averge violation magnitude
-        max_viol: maximum violation magnitude
+    Args
+    - key: tuple (info_provided, seed), where seed may be None when info_provided == 'known'
+    - pkl: saved results read from pickle
+    - ax: optional axes to plot a histogram of violations
+
+    Returns
+    - num_mistakes: total number of mistakes
+    - avg_viol: averge violation magnitude
+    - max_viol: maximum violation magnitude
     """
     vs = pkl['vs']
     assert vs.shape == (T, n)
 
-    violates_max = (vs > V_MAX + 0.05)  # shape [T, n]
-    violates_min = (vs < V_MIN - 0.05)  # shape [T, n]
+    violates_max = (vs > V_MAX2 + 0.05)  # shape [T, n]
+    violates_min = (vs < V_MIN2 - 0.05)  # shape [T, n]
     is_mistake = (violates_max.any(axis=1) | violates_min.any(axis=1))  # shape [T]
 
     num_mistakes = is_mistake.sum()
     num_bus_step_violations = violates_max.sum() + violates_min.sum()
 
     all_violations = np.concatenate([
-        vs[violates_max] - V_MAX,
-        V_MIN - vs[violates_min]
+        vs[violates_max] - V_MAX2,
+        V_MIN2 - vs[violates_min]
     ])
 
     if len(all_violations) == 0:
@@ -79,3 +94,187 @@ def calculate_violations(
           f'max viol: {max_viol:.3g}')
 
     return num_mistakes, avg_viol, max_viol
+
+
+def plot_voltages(
+    vpars: np.ndarray, buses: Sequence[int] = BUSES,
+    ylim: tuple[float, float] | None = (Y_MIN, Y_MAX),
+    yticks: tuple[float, ...] | None = YTICKS,
+    plots_dir: str = '', filename: str = '', legend_filename: str = ''
+) -> None:
+    """Plots voltage curves for selected buses.
+
+    Args
+    - vpars: shape [T, n], squared voltage magnitdues (units kV^2)
+    - buses: which buses to plot, 0 = substation
+    - ylim: (ymin, ymax), set to None to use matplotlib defaults
+    - yticks: sequence of tick locations, set to None to use matplotlib defaults
+    - plots_dir: where to save plots
+    - filename: filename (without extension) for saving plot
+    - legend_filename: filename (without extension) for saving legend plot
+    """
+    T = vpars.shape[0]
+    ts = range(T)
+
+    fig, ax = plt.subplots(figsize=(4, 3), dpi=100, tight_layout=True)
+    for i in np.array(buses) - 1:
+        ax.plot(ts, np.sqrt(vpars[:, i]), label=f'{i+1}')
+
+    ax.axhline(V_MIN, ls='--', color='black')
+    ax.axhline(V_MAX, ls='--', color='black')
+    ax.set(ylabel='Voltage (kV)')
+    if ylim is not None:
+        ax.set(ylim=ylim)
+    if yticks is not None:
+        ax.set(yticks=yticks)
+    ax.set(xlabel='time $t$', xlim=(-50, T),
+           xticks=TIME_TICKS, xticklabels=TIME_LABELS)
+
+    if filename != '':
+        assert plots_dir != ''
+        savefig(fig, plots_dir, filename=f'{filename}.pdf')
+        savefig(fig, plots_dir, filename=f'{filename}.png')
+        savefig(fig, plots_dir, filename=f'{filename}.svg')
+
+    if legend_filename != '':
+        assert plots_dir != ''
+        leg = ax.legend(loc='center left', bbox_to_anchor=(1, 0.5), title='bus')
+        fig.canvas.draw()
+        bbox = leg.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
+        savefig(fig, plots_dir, filename=f'{legend_filename}.pdf', bbox_inches=bbox)
+        savefig(fig, plots_dir, filename=f'{legend_filename}.png', bbox_inches=bbox)
+        savefig(fig, plots_dir, filename=f'{legend_filename}.svg', bbox_inches=bbox)
+
+
+def plot_error_and_etahat(
+    pkls_dict: dict[str, dict[str, Any]], plots_dir: str, filename: str,
+    legend_loc: str | None, etamax: float | None = None
+) -> None:
+    """Plots model error on left axis, etahat on right axis.
+
+    Args
+    - pkls_dict: name (str) => results dict loaded from pickle (dict[str, Any])
+    - plots_dir: where to save plots
+    - filename: filename (without extension) for saving plot
+    - legend_loc: one of [None, 'top', 'separate']
+    - etamax: optional upper limit for etahat axis
+    """
+    fig, ax = plt.subplots(figsize=(4, 3), dpi=100, tight_layout=True)
+    axr = ax.twinx()
+    axr.spines['right'].set_visible(True)
+
+    for name, pkl in pkls_dict.items():
+        print(name)
+        dists = pkl['dists']
+        T = pkl['vs'].shape[0]
+        ax.step(list(dists['t']) + [T], list(dists['X_true']) + [dists['X_true'][-1]],
+                where='post', label=name)
+        ax.scatter(0, dists['X_true'][0])
+        if 'η' in dists:
+            axr.step([0] + list(dists['t']) + [T],
+                     [0] + list(dists['η']) + [dists['η'][-1]], ':',
+                     where='post')
+        else:
+            axr.plot([0, T], [8.65, 8.65], ':')
+
+    ax.set_ylabel(r'$||\hat{X}_t - X^\star||_\bigtriangleup$')
+    axr.set_ylabel(r'$\hat\eta$')
+    ax.set(xticks=TIME_TICKS, xticklabels=TIME_LABELS)
+    ax.set(xlabel='time $t$', xlim=(-50, T))
+
+    if etamax is not None:
+        axr.set_ylim(-0.4, etamax)
+
+    if legend_loc == 'top':
+        ax.legend(ncols=2, bbox_to_anchor=(0, 1), loc='lower left')
+
+    savefig(fig, plots_dir, filename=f'{filename}.pdf')
+    savefig(fig, plots_dir, filename=f'{filename}.png')
+    savefig(fig, plots_dir, filename=f'{filename}.svg')
+
+    if legend_loc == 'separate':
+        axr.set_ylabel('')
+        axr.set_yticklabels([])
+        leg = ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
+        fig.canvas.draw()
+        bbox = leg.get_tightbbox().transformed(fig.dpi_scale_trans.inverted())
+        savefig(fig, plots_dir, f'{filename}_legend.pdf', bbox_inches=bbox)
+        savefig(fig, plots_dir, f'{filename}_legend.png', bbox_inches=bbox)
+        savefig(fig, plots_dir, f'{filename}_legend.svg', bbox_inches=bbox)
+
+
+def plot_fill(
+    ax: plt.Axes, values: np.ndarray, color: int, label: str, alpha: bool = False
+) -> None:
+    """
+    Args
+    - ax: axes
+    - values: shape [num_runs, T]
+    - color: int, index into tab20 colors
+        0 = blue, 2 = orange, 4 = green, 7 = purple
+    - label: line label
+    - alpha: whether to use translucent shading
+    """
+    num_runs, T = values.shape
+    ts = range(T)
+    dark = plt.cm.tab20.colors[color]
+    light = plt.cm.tab20.colors[color + 1]
+
+    if num_runs == 1:
+        ax.plot(ts, values[0], color=dark, lw=0.5, label=label)
+    else:
+        mean = values.mean(axis=0)
+        std = values.std(axis=0)
+        ax.plot(ts, mean, color=dark, lw=0.5, label=label)
+        if alpha:
+            ax.fill_between(ts, mean-std, mean+std, color=light, alpha=0.5)
+        else:
+            ax.fill_between(ts, mean-std, mean+std, color=light)
+
+
+def plot_bus(
+    pkls_by_label: dict[str, Sequence[dict[str, Any]]], bus: int, plots_dir: str = '',
+    legend: bool = False, filename_base: str = ''
+) -> None:
+    """
+    Args:
+    - pkls_by_label: label (str) -> list of results dicts
+    - bus: int, where bus 0 = substation
+    - plots_dir: where to save plots
+    - legend: whether to include legend
+    - filename_base: base filename for plots
+    """
+    fig, ax = plt.subplots(figsize=(4, 3), tight_layout=True)
+    for c, (label, pkls) in enumerate(pkls_by_label.items()):
+        # num_runs = len(pkls)
+        vs = np.stack([
+            np.sqrt(data['vs'][:, bus - 1])
+            for data in pkls
+        ], axis=0)
+        # vs = np.zeros([num_runs, T])
+        # for i, data in enumerate(pkls):
+        #     vs[i] = np.sqrt(data['vs'][:, bus - 1])
+        plot_fill(ax, vs, color=c*2, label=label, alpha=True)
+
+    T = vs.shape[1]
+
+    ax.axhline(11.4, ls='--', color='black')
+    ax.axhline(12.6, ls='--', color='black')
+    ax.set(ylabel='Voltage (kV)', ylim=(Y_MIN, Y_MAX), yticks=YTICKS)
+    ax.set(xlabel='time $t$', xlim=(-50, T),
+           xticks=TIME_TICKS, xticklabels=TIME_LABELS)
+
+    if filename_base != '':
+        assert plots_dir != ''
+
+        savefig(fig, plots_dir, filename=f'{filename_base}_bus{bus}.pdf')
+        savefig(fig, plots_dir, filename=f'{filename_base}_bus{bus}.png')
+        savefig(fig, plots_dir, filename=f'{filename_base}_bus{bus}.svg')
+
+        if legend:
+            leg = ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
+            fig.canvas.draw()
+            bbox = leg.get_tightbbox().transformed(fig.dpi_scale_trans.inverted())
+            savefig(fig, plots_dir, f'{filename_base}_legend.pdf', bbox_inches=bbox)
+            savefig(fig, plots_dir, f'{filename_base}_legend.png', bbox_inches=bbox)
+            savefig(fig, plots_dir, f'{filename_base}_legend.svg', bbox_inches=bbox)