BBH pipeline: write eccentricity params to file, add CLI to compute params & plot

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,

support/Pipelines/Bbh/EccentricityControl.py

@@ -27,6 +27,7 @@ def eccentricity_control(
     tmin: Optional[float] = 500,
     tmax: Optional[float] = None,
     plot_output_dir: Optional[Union[str, Path]] = None,
+    ecc_params_output_file: Optional[Union[str, Path]] = None,
     # Scheduler options
     evolve: bool = True,
     **scheduler_kwargs,
@@ -48,8 +49,6 @@ def eccentricity_control(
       'eccentricity_control_params' in
       'spectre.Pipelines.EccentricityControl.EccentricityControl'.
 
-    - Displays the fit results in a tabular format using a pandas DataFrame.
-
     - If the eccentricity is below a threshold, it prints "Success" and
       indicates that the simulation can continue.
 
@@ -77,42 +76,17 @@ def eccentricity_control(
 
     # Find the current eccentricity and determine new parameters to put into
     # generate-id
-    (
-        eccentricity,
-        ecc_std_dev,
-        new_orbital_params,
-    ) = eccentricity_control_params(
+    ecc_params = eccentricity_control_params(
         h5_files,
         id_input_file_path,
         tmin=tmin,
         tmax=tmax,
         plot_output_dir=plot_output_dir,
+        ecc_params_output_file=ecc_params_output_file,
     )
 
-    # Create DataFrame to display data in tabular format
-    data = {
-        "Attribute": [
-            "Eccentricity",
-            "Eccentricity error",
-            "Updated Omega0",
-            "Updated adot0",
-        ],
-        "Value": [
-            eccentricity,
-            ecc_std_dev,
-            new_orbital_params["Omega0"],
-            new_orbital_params["adot0"],
-        ],
-    }
-    df = pd.DataFrame(data)
-    # Print header line
-    print("=" * 40)
-    # Display table
-    print(df.to_string(index=False))
-    print("=" * 40)
-
     # Stop eccentricity control if eccentricity is below threshold
-    if eccentricity < 0.001:
+    if ecc_params["Eccentricity"] < 0.001:
         print("Success")
         # Should continue the simulation either by restarting from a
         # checkpoint, or from the volume data - will do later
@@ -129,8 +103,8 @@ def eccentricity_control(
         target_params,
         # New orbital parameters
         separation=separation,
-        orbital_angular_velocity=new_orbital_params["Omega0"],
-        radial_expansion_velocity=new_orbital_params["adot0"],
+        orbital_angular_velocity=ecc_params["NewOmega0"],
+        radial_expansion_velocity=ecc_params["NewAdot0"],
         # Initial guesses for ID control
         conformal_mass_a=binary_data["ObjectRight"]["KerrSchild"]["Mass"],
         conformal_mass_b=binary_data["ObjectLeft"]["KerrSchild"]["Mass"],
@@ -163,6 +137,7 @@ def eccentricity_control(
 @click.option(
     "--evolve/--no-evolve",
     default=True,
+    show_default=True,
     help=(
         "Evolve the initial data after generation to continue eccentricity "
         "control. You can disable this to generate only the new initial data "

support/Pipelines/Bbh/Inspiral.yaml

@@ -21,8 +21,9 @@ Next:
   Run: spectre.Pipelines.Bbh.EccentricityControl:eccentricity_control
   With:
     h5_files: "../Segment_*/BbhReductions.h5"
-    plot_output_dir: ./
     id_input_file_path: {{id_input_file_path }}
+    plot_output_dir: ./
+    ecc_params_output_file: "../EccentricityParams.yaml"
     pipeline_dir: {{ pipeline_dir }}
     scheduler: {{ scheduler | default("None") }}
     copy_executable: {{ copy_executable | default("None") }}