-
Notifications
You must be signed in to change notification settings - Fork 189
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #5929 from nilsvu/plot_ellsolver_convergence
Add CLI command to plot elliptic solver convergence
- Loading branch information
Showing
5 changed files
with
226 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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"]) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
70 changes: 70 additions & 0 deletions
70
tests/Unit/Visualization/Python/Test_PlotEllipticConvergence.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |