# Copyright (c) 2012-2024, OpenGeoSys Community (http://www.opengeosys.org)
# Distributed under a Modified BSD License.
# See accompanying file LICENSE.txt or
# http://www.opengeosys.org/project/license
#
"""Utilities to create nicely spaced levels."""
from math import nextafter
from typing import Any
import numpy as np
from pyvista import UnstructuredGrid
from ogstools.plot.shared import setup
from ogstools.variables import Variable, get_preset
[docs]
def nice_num(val: float) -> float:
"""
Return the closest number of the form 10**x * {1,2,4,5}.
Fractions containing only these number are ensured to have
terminating decimal representations.
"""
pow10 = 10 ** np.floor(np.log10(val))
vals = np.array([1.0, 2.0, 4.0, 5.0, 10.0])
return pow10 * vals[np.argmin(np.abs(val / pow10 - vals))]
[docs]
def nice_range(lower: float, upper: float, n_ticks: float) -> np.ndarray:
"""
Return an array in the interval (lower, upper) with terminating decimals.
The length of the arrays will be close to n_ticks.
"""
base = nice_num(upper - lower)
tick_spacing = nice_num(base / (n_ticks - 1))
nice_lower = np.ceil(lower / tick_spacing) * tick_spacing
nice_upper = np.ceil(upper / tick_spacing) * tick_spacing
res = np.arange(nice_lower, nice_upper, tick_spacing)
return res[(res > lower) & (res < upper)]
[docs]
def adaptive_rounding(vals: np.ndarray, precision: int) -> np.ndarray:
"""
Return the given values rounded to significant digits.
The significant digits are based of the median decimal exponent and the
given precision.
"""
if vals.size == 0:
return vals
median_exp = median_exponent(vals)
rounded_vals = np.stack([np.round(v, precision - median_exp) for v in vals])
if len(set(rounded_vals)) > 1:
return rounded_vals
return np.stack([np.round(v, 12 - median_exp) for v in vals])
[docs]
def compute_levels(lower: float, upper: float, n_ticks: int) -> np.ndarray:
"""
Return an array in the interval [lower, upper] with terminating decimals.
The length of the arrays will be close to n_ticks.
At the boundaries the tickspacing may differ from the remaining array.
"""
if lower == upper:
return np.asarray([lower, nextafter(lower, np.inf)])
result = nice_range(lower, upper, n_ticks)
levels = np.unique(
adaptive_rounding(
np.append(np.append(lower, result), upper), precision=3
)
)
if levels[0] == levels[-1]:
return np.array([levels[0], nextafter(levels[0], np.inf)])
return levels
[docs]
def combined_levels(
meshes: np.ndarray, variable: Variable | str, **kwargs: Any
) -> np.ndarray:
"""
Calculate well spaced levels for the encompassing variable range in meshes.
"""
variable = get_preset(variable, meshes.ravel()[0])
vmin, vmax = np.inf, -np.inf
VMIN = kwargs.get("vmin", setup.vmin)
VMAX = kwargs.get("vmax", setup.vmax)
unique_vals = np.array([])
mesh: UnstructuredGrid
for mesh in np.ravel(meshes):
values = variable.magnitude.transform(
mesh.ctp(True).threshold(value=[1, 1], scalars=variable.mask)
if variable.mask_used(mesh)
else mesh
)
if (
kwargs.get("log_scaled", setup.log_scaled)
and not variable.is_mask()
):
values = np.log10(np.where(values > 1e-14, values, 1e-14))
vmin = min(vmin, np.nanmin(values)) if VMIN is None else vmin
vmax = max(vmax, np.nanmax(values)) if VMAX is None else vmax
unique_vals = np.unique(
np.concatenate((unique_vals, np.unique(values)))
)
vmin = vmin if VMIN is None else VMIN
vmax = vmax if VMAX is None else VMAX
if vmin == vmax:
return np.array([vmin, nextafter(vmax, np.inf)])
if (
all(val.is_integer() for val in unique_vals)
and VMIN is None
and VMAX is None
):
return unique_vals[(vmin <= unique_vals) & (unique_vals <= vmax)]
return compute_levels(
vmin, vmax, kwargs.get("num_levels", setup.num_levels)
)
[docs]
def level_boundaries(levels: np.ndarray) -> np.ndarray:
return np.array(
[
levels[0] - 0.5 * (levels[1] - levels[0]),
*0.5 * (levels[:-1] + levels[1:]),
levels[-1] + 0.5 * (levels[-1] - levels[-2]),
]
)
