From 23893851b0c4abc81843404a47a3325afb202609 Mon Sep 17 00:00:00 2001 From: Nils Vu Date: Fri, 19 Apr 2024 12:35:00 -0700 Subject: [PATCH] Add CLI command to plot elliptic solver convergence --- src/Visualization/Python/CMakeLists.txt | 1 + .../Python/PlotEllipticConvergence.py | 137 ++++++++++++++++++ src/Visualization/Python/__init__.py | 7 + .../Unit/Visualization/Python/CMakeLists.txt | 11 ++ .../Python/Test_PlotEllipticConvergence.py | 70 +++++++++ 5 files changed, 226 insertions(+) create mode 100644 src/Visualization/Python/PlotEllipticConvergence.py create mode 100644 tests/Unit/Visualization/Python/Test_PlotEllipticConvergence.py diff --git a/src/Visualization/Python/CMakeLists.txt b/src/Visualization/Python/CMakeLists.txt index 9e129da2a6f0..acd57ec54a63 100644 --- a/src/Visualization/Python/CMakeLists.txt +++ b/src/Visualization/Python/CMakeLists.txt @@ -13,6 +13,7 @@ spectre_python_add_module( Plot.py PlotAlongLine.py PlotDatFile.py + PlotEllipticConvergence.py PlotMemoryMonitors.py PlotPowerMonitors.py PlotSizeControl.py diff --git a/src/Visualization/Python/PlotEllipticConvergence.py b/src/Visualization/Python/PlotEllipticConvergence.py new file mode 100644 index 000000000000..0a0e23daaec0 --- /dev/null +++ b/src/Visualization/Python/PlotEllipticConvergence.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python + +# Distributed under the MIT License. +# See LICENSE.txt for details. + +import logging +from typing import List + +import click +import h5py +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from matplotlib.ticker import MaxNLocator + +from spectre.Visualization.Plot import ( + apply_stylesheet_command, + show_or_save_plot_command, +) +from spectre.Visualization.ReadH5 import to_dataframe + +logger = logging.getLogger(__name__) + + +def split_iteration_sequence(data: pd.DataFrame) -> List[pd.DataFrame]: + """Split a dataframe where its index is not strictly increasing. + + Returns a list of dataframes, where each has a strictly increasing index. + """ + split_indices = np.split( + np.arange(len(data)), np.nonzero(np.diff(data.index) < 1)[0] + 1 + ) + return [data.iloc[split_index] for split_index in split_indices] + + +def plot_elliptic_convergence( + h5_file, + ax=None, + linear_residuals_subfile_name="GmresResiduals.dat", + nonlinear_residuals_subfile_name="NewtonRaphsonResiduals.dat", +): + """Plot elliptic solver convergence + + Arguments: + h5_file: The H5 reductions file. + ax: Optional. The matplotlib axis to plot on. + linear_residuals_subfile_name: The name of the subfile containing the + linear solver residuals. + nonlinear_residuals_subfile_name: The name of the subfile containing the + nonlinear solver residuals. + """ + with h5py.File(h5_file, "r") as open_h5file: + linear_residuals = split_iteration_sequence( + to_dataframe(open_h5file[linear_residuals_subfile_name]).set_index( + "Iteration" + ) + ) + nonlinear_residuals = ( + split_iteration_sequence( + to_dataframe( + open_h5file[nonlinear_residuals_subfile_name] + ).set_index("Iteration") + ) + if nonlinear_residuals_subfile_name in open_h5file + else None + ) + cumulative_linsolv_iterations = [0] + list( + np.cumsum([len(l) - 1 for l in linear_residuals]) + ) + norm = ( + nonlinear_residuals + if nonlinear_residuals is not None + else linear_residuals + )[0]["Residual"].iloc[0] + # Plot nonlinear solver residuals + if ax is not None: + plt.sca(ax) + if nonlinear_residuals is not None: + m = 0 + for i, residuals in enumerate(nonlinear_residuals): + plt.plot( + cumulative_linsolv_iterations[m : m + len(residuals)], + residuals["Residual"] / norm, + color="black", + ls="dotted", + marker=".", + label="Nonlinear residual" if i == 0 else None, + ) + m += len(residuals) - 1 + # Plot linear solver residuals + for i, residuals in enumerate(linear_residuals): + plt.plot( + residuals.index + cumulative_linsolv_iterations[i], + residuals["Residual"] / norm, + color="black", + label="Linear residual" if i == 0 else None, + marker="." if len(residuals) < 20 else None, + ) + + # Configure the axes + plt.yscale("log") + plt.grid() + plt.legend() + plt.xlabel("Cumulative linear solver iteration") + plt.ylabel("Relative residual") + plt.title("Elliptic solver convergence") + # Allow only integer ticks for the x-axis + plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True)) + + +@click.command(name="elliptic-convergence") +@click.argument( + "h5_file", + type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True), +) +@click.option( + "--linear-residuals-subfile-name", + help="The name of the subfile containing the linear solver residuals", + default="GmresResiduals.dat", + show_default=True, +) +@click.option( + "--nonlinear-residuals-subfile-name", + help="The name of the subfile containing the nonlinear solver residuals", + default="NewtonRaphsonResiduals.dat", + show_default=True, +) +@apply_stylesheet_command() +@show_or_save_plot_command() +def plot_elliptic_convergence_command(**kwargs): + """Plot elliptic solver convergence""" + _rich_traceback_guard = True # Hide traceback until here + plot_elliptic_convergence(**kwargs) + + +if __name__ == "__main__": + plot_elliptic_convergence_command(help_option_names=["-h", "--help"]) diff --git a/src/Visualization/Python/__init__.py b/src/Visualization/Python/__init__.py index 5440c0df66a8..225945a685f6 100644 --- a/src/Visualization/Python/__init__.py +++ b/src/Visualization/Python/__init__.py @@ -9,6 +9,7 @@ def list_commands(self, ctx): return [ "along-line", "dat", + "elliptic-convergence", "memory-monitors", "power-monitors", "size-control", @@ -26,6 +27,12 @@ def get_command(self, ctx, name): from spectre.Visualization.PlotDatFile import plot_dat_command return plot_dat_command + elif name == "elliptic-convergence": + from spectre.Visualization.PlotEllipticConvergence import ( + plot_elliptic_convergence_command, + ) + + return plot_elliptic_convergence_command elif name == "memory-monitors": from spectre.Visualization.PlotMemoryMonitors import ( plot_memory_monitors_command, diff --git a/tests/Unit/Visualization/Python/CMakeLists.txt b/tests/Unit/Visualization/Python/CMakeLists.txt index dbbefc90949f..8c4c5de79745 100644 --- a/tests/Unit/Visualization/Python/CMakeLists.txt +++ b/tests/Unit/Visualization/Python/CMakeLists.txt @@ -19,6 +19,12 @@ spectre_add_python_bindings_test( "unit;visualization;python" None) +spectre_add_python_bindings_test( + "Unit.Visualization.Python.PlotEllipticConvergence" + Test_PlotEllipticConvergence.py + "unit;visualization;python" + None) + spectre_add_python_bindings_test( "Unit.Visualization.Python.PlotMemoryMonitors" Test_PlotMemoryMonitors.py @@ -56,6 +62,11 @@ if(${BUILD_PYTHON_BINDINGS}) PROPERTIES TIMEOUT 10 ) + set_tests_properties( + "Unit.Visualization.Python.PlotEllipticConvergence" + PROPERTIES + TIMEOUT 10 + ) set_tests_properties( "Unit.Visualization.Python.PlotMemoryMonitors" PROPERTIES diff --git a/tests/Unit/Visualization/Python/Test_PlotEllipticConvergence.py b/tests/Unit/Visualization/Python/Test_PlotEllipticConvergence.py new file mode 100644 index 000000000000..dd765b1ab7ca --- /dev/null +++ b/tests/Unit/Visualization/Python/Test_PlotEllipticConvergence.py @@ -0,0 +1,70 @@ +# Distributed under the MIT License. +# See LICENSE.txt for details. + +import os +import shutil +import unittest + +import click +import numpy as np +from click.testing import CliRunner + +import spectre.Informer as spectre_informer +import spectre.IO.H5 as spectre_h5 +from spectre.Visualization.PlotEllipticConvergence import ( + plot_elliptic_convergence_command, +) + + +class TestPlotEllipticConvergence(unittest.TestCase): + def setUp(self): + self.test_dir = os.path.join( + spectre_informer.unit_test_build_path(), + "Visualization/PlotEllipticConvergence", + ) + shutil.rmtree(self.test_dir, ignore_errors=True) + self.h5_filename = os.path.join(self.test_dir, "Residuals.h5") + os.makedirs(self.test_dir, exist_ok=True) + with spectre_h5.H5File(self.h5_filename, "w") as h5_file: + linear_residuals = h5_file.insert_dat( + "/GmresResiduals", + legend=["Iteration", "Residual"], + version=0, + ) + linear_residuals.append([0, 1.0]) + linear_residuals.append([1, 0.5]) + linear_residuals.append([0, 0.6]) + linear_residuals.append([1, 0.4]) + linear_residuals.append([2, 0.2]) + h5_file.close_current_object() + nonlinear_residuals = h5_file.insert_dat( + "/NewtonRaphsonResiduals", + legend=["Iteration", "Residual"], + version=0, + ) + nonlinear_residuals.append([0, 1.0]) + nonlinear_residuals.append([1, 0.5]) + + def tearDown(self): + shutil.rmtree(self.test_dir) + + def test_cli(self): + output_filename = os.path.join(self.test_dir, "output.pdf") + runner = CliRunner() + result = runner.invoke( + plot_elliptic_convergence_command, + [ + self.h5_filename, + "-o", + output_filename, + ], + catch_exceptions=False, + ) + self.assertEqual(result.exit_code, 0) + # We can't test that the plot is correct, so just make sure it at least + # gets written out. + self.assertTrue(os.path.exists(output_filename)) + + +if __name__ == "__main__": + unittest.main(verbosity=2)