"""Density field plot operator."""
from __future__ import annotations
import matplotlib.axes
import matplotlib.colors
import numpy as np
from tud_lbm.io.plotting.base import PlotOperator
from tud_lbm.registry import plotting_operator
# Density ratio threshold for using log scale
[docs]
DENSITY_RATIO_THRESHOLD = 100
@plotting_operator(name="density")
[docs]
class DensityPlotOperator(PlotOperator):
"""Render the density field as a 2-D colour map."""
[docs]
def is_available(self, data: dict[str, np.ndarray]) -> bool:
"""Check if density data is available for rendering."""
return "rho" in data
[docs]
def __call__(
self,
ax: matplotlib.axes.Axes,
data: dict[str, np.ndarray],
timestep: int,
) -> None:
"""Render the density field as a 2D color map."""
rho = np.asarray(data["rho"])[:, :, 0, 0, 0].T
use_log_scale = False
if "multiphase" in self.config.sim_type:
rho_l = self.config.rho_l
rho_v = self.config.rho_v
if rho_l and rho_v and rho_v != 0:
use_log_scale = (rho_l / rho_v) > DENSITY_RATIO_THRESHOLD
if use_log_scale:
positive = rho[rho > 0]
floor = float(np.min(positive) * 1e-10) if positive.size else 1e-30
plot_data = np.maximum(rho, floor)
im = ax.imshow(
plot_data,
origin="lower",
aspect="equal",
cmap="viridis",
norm=matplotlib.colors.LogNorm(),
)
ax.set_title(f"Density (log) t={timestep}")
else:
im = ax.imshow(rho, origin="lower", aspect="equal", cmap="viridis")
ax.set_title(f"Density t={timestep}")
ax.figure.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="Density")
ax.set_xlabel("x")
ax.set_ylabel("y")