Add CLI to plot ecc control iterations

src/Visualization/Python/CMakeLists.txt

@@ -16,6 +16,7 @@ spectre_python_add_module(
   PlotCce.py
   PlotControlSystem.py
   PlotDatFile.py
+  PlotEccentricityControl.py
   PlotEllipticConvergence.py
   PlotMemoryMonitors.py
   PlotPowerMonitors.py

src/Visualization/Python/PlotEccentricityControl.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python
+
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+import logging
+from pathlib import Path
+from typing import Sequence, Union
+
+import click
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import yaml
+from matplotlib.colors import LogNorm
+from scipy.interpolate import CloughTocher2DInterpolator
+
+from spectre.Visualization.Plot import (
+    apply_stylesheet_command,
+    show_or_save_plot_command,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def plot_eccentricity_control(
+    ecc_params_files: Sequence[Union[str, Path]],
+):
+    """Plot eccentricity control iterations.
+
+    \f
+    Arguments:
+      ecc_params_files: YAML files containing the output of
+        'eccentricity_control_params()'.
+
+    Returns:
+      The figure containing the plot.
+    """
+    # Load the eccentricity parameters
+    ecc_params = []
+    for ecc_params_file in ecc_params_files:
+        with open(ecc_params_file, "r") as open_file:
+            ecc_params.append(yaml.safe_load(open_file))
+    ecc_params = pd.DataFrame(ecc_params)
+
+    # Set up the figure
+    fig = plt.figure(figsize=(4.5, 4))
+    plt.xlabel(r"Angular velocity $\Omega_0$")
+    plt.ylabel(r"Expansion $\dot{a}_0$")
+    plt.grid(zorder=0)
+    plt.ticklabel_format(axis="both", style="sci", scilimits=(0, 0))
+
+    # Plot measured eccentricity contours
+    ecc_norm = LogNorm(vmin=1e-4, vmax=1)
+    if len(ecc_params) > 2:
+        ecc_intrp = CloughTocher2DInterpolator(
+            ecc_params[["Omega0", "Adot0"]],
+            np.log10(ecc_params["Eccentricity"]),
+            rescale=True,
+        )
+        x, y = np.meshgrid(
+            np.linspace(*ecc_params["Omega0"].agg(["min", "max"]), 200),
+            np.linspace(*ecc_params["Adot0"].agg(["min", "max"]), 200),
+        )
+        z = ecc_intrp(x, y)
+        ecc_contours = plt.contourf(
+            x,
+            y,
+            10**z,
+            levels=10.0 ** np.arange(-4, 1),
+            norm=ecc_norm,
+            cmap="Blues",
+            zorder=10,
+        )
+    else:
+        ecc_contours = plt.cm.ScalarMappable(
+            norm=ecc_norm,
+            cmap="Blues",
+        )
+    plt.colorbar(ecc_contours, label="Eccentricity", ax=plt.gca())
+    cbar_ax = fig.axes[1]
+    plt.gca().use_sticky_edges = False
+
+    # Plot the path through parameter space
+    plt.plot(
+        ecc_params["Omega0"],
+        ecc_params["Adot0"],
+        color="black",
+        zorder=20,
+    )
+
+    # Place a marker at each iteration
+    scatter_kwargs = dict(
+        color="white",
+        marker="o",
+        linewidth=1,
+        edgecolor="black",
+        s=100,
+        zorder=25,
+    )
+    annotate_kwargs = dict(
+        textcoords="offset points",
+        xytext=(0, -0.5),
+        ha="center",
+        va="center",
+        fontsize=8,
+        zorder=30,
+    )
+    plt.scatter(ecc_params["Omega0"], ecc_params["Adot0"], **scatter_kwargs)
+    cbar_ax.scatter(
+        np.repeat(0.5, len(ecc_params)),
+        ecc_params["Eccentricity"],
+        **scatter_kwargs,
+    )
+    for i in range(len(ecc_params)):
+        plt.annotate(
+            f"{i}",
+            ecc_params.iloc[i][["Omega0", "Adot0"]],
+            **annotate_kwargs,
+        )
+        cbar_ax.annotate(
+            f"{i}", (0.5, ecc_params["Eccentricity"][i]), **annotate_kwargs
+        )
+
+    # Plot location of the next iteration
+    plt.plot(
+        ecc_params.iloc[-1][["Omega0", "NewOmega0"]],
+        ecc_params.iloc[-1][["Adot0", "NewAdot0"]],
+        color="black",
+        linestyle="dotted",
+        zorder=20,
+    )
+    plt.scatter(
+        ecc_params.iloc[-1]["NewOmega0"],
+        ecc_params.iloc[-1]["NewAdot0"],
+        color="black",
+        marker="o",
+        s=30,
+        zorder=40,
+    )
+    return fig
+
+
+@click.command(
+    name="eccentricity-control", help=plot_eccentricity_control.__doc__
+)
+@click.argument(
+    "ecc_params_files",
+    nargs=-1,
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True),
+)
+@apply_stylesheet_command()
+@show_or_save_plot_command()
+def plot_eccentricity_control_command(**kwargs):
+    return plot_eccentricity_control(**kwargs)
+
+
+if __name__ == "__main__":
+    plot_eccentricity_control_command(help_option_names=["-h", "--help"])

src/Visualization/Python/__init__.py

@@ -13,6 +13,7 @@ def list_commands(self, ctx):
             "cce",
             "control-system",
             "dat",
+            "eccentricity-control",
             "elliptic-convergence",
             "memory-monitors",
             "power-monitors",
@@ -42,6 +43,12 @@ def get_command(self, ctx, name):
             from spectre.Visualization.PlotDatFile import plot_dat_command
 
             return plot_dat_command
+        elif name in ["eccentricity-control", "ecc-control"]:
+            from spectre.Visualization.PlotEccentricityControl import (
+                plot_eccentricity_control_command,
+            )
+
+            return plot_eccentricity_control_command
         elif name == "elliptic-convergence":
             from spectre.Visualization.PlotEllipticConvergence import (
                 plot_elliptic_convergence_command,

tests/Unit/Visualization/Python/CMakeLists.txt

@@ -41,6 +41,13 @@ spectre_add_python_bindings_test(
   None
   TIMEOUT 10)
 
+spectre_add_python_bindings_test(
+  "Unit.Visualization.Python.PlotEccentricityControl"
+  Test_PlotEccentricityControl.py
+  "unit;visualization;python"
+  None
+  TIMEOUT 10)
+
 spectre_add_python_bindings_test(
   "Unit.Visualization.Python.PlotEllipticConvergence"
   Test_PlotEllipticConvergence.py

tests/Unit/Visualization/Python/Test_PlotEccentricityControl.py

@@ -0,0 +1,59 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+import os
+import shutil
+import unittest
+
+import numpy as np
+import yaml
+from click.testing import CliRunner
+
+from spectre.Informer import unit_test_build_path
+from spectre.Visualization.PlotEccentricityControl import (
+    plot_eccentricity_control_command,
+)
+
+
+class TestPlotEccentricityControl(unittest.TestCase):
+    def setUp(self):
+        self.test_dir = os.path.join(
+            unit_test_build_path(), "Visualization", "PlotEccentricityControl"
+        )
+        os.makedirs(self.test_dir, exist_ok=True)
+        self.ecc_filenames = []
+        for i in range(3):
+            ecc_filename = os.path.join(
+                self.test_dir,
+                f"EccentricityParams{i}.yaml",
+            )
+            with open(ecc_filename, "w") as open_file:
+                yaml.safe_dump(
+                    {
+                        "Eccentricity": 10 ** (-i),
+                        "Omega0": 0.015 + np.random.rand() * 1e-4,
+                        "Adot0": 1e-4 + np.random.rand() * 1e-4,
+                        "NewOmega0": 0.015 + np.random.rand() * 1e-4,
+                        "NewAdot0": 1e-4 + np.random.rand() * 1e-4,
+                    },
+                    open_file,
+                )
+            self.ecc_filenames.append(ecc_filename)
+
+    def tearDown(self):
+        shutil.rmtree(self.test_dir)
+
+    def test_cli(self):
+        runner = CliRunner()
+        output_filename = os.path.join(self.test_dir, "output.pdf")
+        result = runner.invoke(
+            plot_eccentricity_control_command,
+            self.ecc_filenames + ["-o", output_filename],
+            catch_exceptions=False,
+        )
+        self.assertEqual(result.exit_code, 0, result.output)
+        self.assertTrue(os.path.exists(output_filename))
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)