# SPDX-FileCopyrightText: Copyright (c) OpenGeoSys Community (opengeosys.org)
# SPDX-License-Identifier: BSD-3-Clause
from typing import Any, Literal
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.ticker import MaxNLocator
from ogstools.plot import heatmap
from ogstools.variables import Scalar
from .common_ogs_analyses import (
convergence_order_per_ts_iteration,
errors_per_ts_iteration,
model_and_clock_time,
)
def _format_fig(
fig: plt.Figure,
ax: plt.Axes,
x_ticks: np.ndarray,
y_ticks: np.ndarray,
x_label: str,
) -> plt.Figure:
x_axis_labels = {
"time_step": "time step",
"model_time": "model time [s]",
}
ax.set_xticks(x_ticks, minor=True)
ax.set_yticks(y_ticks, minor=True)
ax.yaxis.set_major_locator(MaxNLocator(integer=True))
ax.set_xlabel(x_axis_labels.get(x_label, x_label.replace("_", " ")))
ax.set_ylabel("iteration")
fig.tight_layout()
return fig
def _x_data(
df: pd.DataFrame,
x_metric: Literal["time_step", "model_time"],
default_len_x: int,
) -> tuple[np.ndarray | None, range]:
if x_metric == "time_step":
return None, range(default_len_x)
df_time = model_and_clock_time(df).reset_index()
df_time = df_time.pivot_table(
"model_time", "time_step", aggfunc="last"
).reset_index()
x_vals = df_time[x_metric].to_numpy()
x_ticks = 0.5 * (x_vals[1:] + x_vals[:-1])
if len(x_vals) == default_len_x:
# TODO: Not sure why, but sometimes the required length of the x_vals
# is off by 1 compared to the required length to plot a heatmap with
# errors or orders. The following fixes this but the underlying problem
# is not yet understood.
x_vals = np.append(0.0, x_vals)
return x_vals, x_ticks
[docs]
def plot_convergence_order(
df: pd.DataFrame,
n: Literal[3, 4] = 3,
x_metric: Literal["time_step", "model_time"] = "time_step",
**kwargs: Any,
) -> plt.Figure:
"""
Create a heatmap of the nonlinear solver convergence order.
see: :func:`~ogstools.logparser.common_ogs_analyses.convergence_order_per_ts_iteration`
In order to estimate the convergence order we need to take into account
multiple values and thus cannot assign each iteration a convergence order.
Only for iterations `i` of `i >= n` an order is calculated and plotted.
Per default the scale is limited to a range of 0 to 2 to limit the view to
meaningful data. Set the keyword arguments `vmin` and `vmax`to `None` to
see the entire scale.
:param df: Dataframe of a simulation log.
:param n: Number of error values to use to estimate the
convergence order.
:param x_metric: x_axis can represent either "time_step" or
"model_time" (not computation time)
Keyword Arguments:
- see: :func:`~ogstools.plot.heatmaps.heatmap`)
:returns: A figure with a heatmap of the nonlinear solver convergence order.
"""
orders = convergence_order_per_ts_iteration(df, n=n)
x_vals, x_ticks = _x_data(df, x_metric, orders.shape[1])
y_ticks = range(orders.shape[0])
order_var = Scalar("convergence_order", cmap="RdBu", symbol="q")
kwargs.setdefault("vmin", 0)
kwargs.setdefault("vmax", 2)
res = heatmap(orders, order_var, x_vals=x_vals, **kwargs)
fig: plt.Figure = kwargs.get("fig", res)
ax = kwargs.get("ax", fig.axes[0])
return _format_fig(fig, ax, x_ticks, y_ticks, x_label=x_metric)
[docs]
def plot_convergence(
df: pd.DataFrame,
metric: Literal["dx", "dx_x", "x"] = "dx",
x_metric: Literal["time_step", "model_time"] = "time_step",
**kwargs: Any,
) -> plt.Figure:
"""
Create a heatmap of the nonlinear solver convergence data.
The individual values in the heatmap correspond to the top right indices on
the x- and y-axis. E.g. the very first entry which fills the space between
timesteps 0-1 and iteration 0-1 belongs to the first iteration of the first
timestep. Thus we immediately read on which iteration a timestep converged
and on which timestep the simulation ended. Per default logarithmic scaling
is used. Set `log_scaled` to `False` to use linear scaling.
:param df: Dataframe of a simulation log.
:param metric: Which metric / column of the Dataframe to plot.
dx (absolute error), dx_x (relative error), x (residual)
:param x_metric: x_axis can represent either "time_step" or
"model_time" (not computation time)
Keyword Arguments:
- see: :func:`~ogstools.plot.heatmaps.heatmap`)
:returns: A figure with a heatmap of the nonlinear solver convergence data.
"""
errors = errors_per_ts_iteration(df, metric)
x_vals, x_ticks = _x_data(df, x_metric, errors.shape[1])
y_ticks = range(errors.shape[0])
names = {"dx": "absolute error", "dx_x": "relative error", "x": "residual"}
symbol = str(metric).replace("_", " / ")
err_var = Scalar(names[metric], cmap="viridis", symbol=symbol)
kwargs.setdefault("log_scaled", True)
res = heatmap(errors, err_var, x_vals=x_vals, **kwargs)
fig: plt.Figure = kwargs.get("fig", res)
ax = kwargs.get("ax", fig.axes[0])
return _format_fig(fig, ax, x_ticks, y_ticks, x_label=x_metric)
