"""
Standalone extraction of 2D and 3D simulation results at user points.
This module exposes the depth-explicit extraction that drives the calibration
workflow (``TelemacModel.extract_data_point``) as a plain function that works
on any TELEMAC SELAFIN file (2D or 3D) or OpenFOAM VTK output without a model
instance, configuration file, or results-folder conventions::
import hydroBayesCal as hbc
df = hbc.extract_results("r3d_case.slf", variable="velocity",
x=[5.0, 7.5], y=[2.0, 2.5], z=0.05)
For 3D TELEMAC results the vertical plane closest to the requested height
above the local bed is selected via ``ELEVATION Z`` (same convention as the
calibration workflow). With ``z=None`` the full vertical profile is returned.
2D results and the 2D files generated by TELEMAC-3D are handled by the same
function. OpenFOAM VTK output (``foamToVTK``) is inherently 3D and sampled at
the exact (x, y, z) locations.
"""
import glob
import os
import numpy as np
import pandas as pd
from scipy import spatial
from hydroBayesCal.telemac.pputils.ppmodules.selafin_io_pp import ppSELAFIN
#: File endings treated as TELEMAC SELAFIN files.
SELAFIN_EXTENSIONS = (".slf", ".sel", ".srf", ".res")
#: File endings treated as OpenFOAM/VTK files.
VTK_EXTENSIONS = (".vtu", ".vtk", ".vtm", ".vtp")
#: Convenience aliases resolved against the variables present in the file.
TELEMAC_ALIASES = {
"DEPTH": "WATER DEPTH",
"WSE": "FREE SURFACE",
"WATER SURFACE": "FREE SURFACE",
"BED": "BOTTOM",
}
def _reduce_time(all_times_values, time, time_index, n_last):
"""Reduce a (n_times, ...) array to one time step according to the mode."""
all_times_values = np.asarray(all_times_values, dtype=float)
if time == "last":
return all_times_values[-1]
if time == "index":
return all_times_values[time_index]
return np.mean(all_times_values[-n_last:], axis=0)
def _idw_weights(node_xy, target_xy):
"""Inverse-distance weights of nodes at node_xy relative to target_xy."""
distances = np.linalg.norm(
np.asarray(node_xy, dtype=float) - np.asarray(target_xy, dtype=float),
axis=1
)
distances[distances == 0] = np.finfo(float).eps
weights = 1.0 / distances
return weights / weights.sum()
# ======================================================================
# TELEMAC (SELAFIN)
# ======================================================================
def _resolve_telemac_variables(requested, var_index, is_3d):
"""
Map requested variable names to columns of the SLF variable table.
Returns a list of (requested_name, component_indices) where the value is
the plain variable column for single indices, or the vector magnitude for
multiple indices.
"""
resolved = []
for name in requested:
key = " ".join(str(name).split()).upper()
key = TELEMAC_ALIASES.get(key, key)
if key in var_index:
resolved.append((name, [var_index[key]]))
continue
if key in ("VELOCITY", "3D VELOCITY MAGNITUDE", "SCALAR VELOCITY"):
if not is_3d and "SCALAR VELOCITY" in var_index:
resolved.append((name, [var_index["SCALAR VELOCITY"]]))
continue
components = ["VELOCITY U", "VELOCITY V"]
if is_3d and "VELOCITY W" in var_index:
components.append("VELOCITY W")
missing = [c for c in components if c not in var_index]
if missing:
raise ValueError(
f"Cannot resolve '{name}': the file provides neither "
f"SCALAR VELOCITY nor the components {missing}. "
f"Available variables: {sorted(var_index)}"
)
resolved.append((name, [var_index[c] for c in components]))
continue
raise ValueError(
f"Variable '{name}' not found in the SELAFIN file. "
f"Available variables: {sorted(var_index)}"
)
return resolved
def _evaluate_variables(resolved_variables, values):
"""Evaluate resolved variables on one row of SLF values."""
out = {}
for name, indices in resolved_variables:
if len(indices) == 1:
out[name] = float(values[indices[0]])
else:
out[name] = float(np.sqrt(np.sum(values[indices] ** 2)))
return out
def _extract_telemac(result_file, variables, x_pts, y_pts, z_pts,
time, time_index, n_last, method, k):
if not os.path.isfile(result_file):
raise FileNotFoundError(f"SELAFIN file not found: {result_file}")
slf = ppSELAFIN(result_file)
slf.readHeader()
slf.readTimes()
var_names = [" ".join(v.split()) for v in slf.getVarNames()]
var_index = {v.upper(): i for i, v in enumerate(var_names)}
n_var = len(var_names)
_, npoin, _, _, _, mesh_x, mesh_y = slf.getMesh()
nplan = max(slf.getNPLAN(), 1)
is_3d = nplan > 1
n2d = npoin // nplan
resolved_variables = _resolve_telemac_variables(variables, var_index, is_3d)
if is_3d and "ELEVATION Z" not in var_index:
raise ValueError(
"The 3D SELAFIN file does not provide 'ELEVATION Z', which is "
"required to locate the vertical planes. "
f"Available variables: {sorted(var_index)}"
)
tree = spatial.cKDTree(np.column_stack((mesh_x[:n2d], mesh_y[:n2d])))
k_use = 1 if method == "nearest" else min(k, n2d)
records = []
for i in range(x_pts.size):
target_xy = (x_pts[i], y_pts[i])
_, idx = tree.query(target_xy, k=k_use)
idx_base = np.atleast_1d(idx)
node_xy = np.column_stack((mesh_x[idx_base], mesh_y[idx_base]))
weights = _idw_weights(node_xy, target_xy) if k_use > 1 else np.array([1.0])
# Time-reduced values of all variables on all planes above the
# selected horizontal nodes: shape (nplan, k_use, n_var).
node_values = np.zeros((nplan, k_use, n_var))
for p in range(nplan):
for j, base in enumerate(idx_base):
slf.readVariablesAtNode(int(base + p * n2d))
node_values[p, j, :] = _reduce_time(
slf.getVarValuesAtNode(), time, time_index, n_last
)
# IDW across the k nodes on every plane: shape (nplan, n_var).
plane_values = np.einsum("j,pjv->pv", weights, node_values)
base_record = {"point": i, "x": x_pts[i], "y": y_pts[i]}
if not is_3d:
records.append(
{**base_record, **_evaluate_variables(resolved_variables,
plane_values[0])}
)
continue
z_profile = plane_values[:, var_index["ELEVATION Z"]]
bed_z = z_profile[0]
if z_pts is not None:
# Select the plane closest to the requested height above the bed
# (same convention as TelemacModel.extract_data_point).
p = int(np.argmin(np.abs(z_profile - (bed_z + z_pts[i]))))
records.append({
**base_record,
"z": z_pts[i],
"plane": p + 1,
"z_plane": z_profile[p],
"z_above_bed": z_profile[p] - bed_z,
**_evaluate_variables(resolved_variables, plane_values[p]),
})
else:
# Full vertical profile: one row per plane.
for p in range(nplan):
records.append({
**base_record,
"plane": p + 1,
"z_plane": z_profile[p],
"z_above_bed": z_profile[p] - bed_z,
**_evaluate_variables(resolved_variables, plane_values[p]),
})
return pd.DataFrame(records)
# ======================================================================
# OpenFOAM (VTK)
# ======================================================================
def _collect_vtk_files(source):
"""Return the time-sorted VTK files of an OpenFOAM case directory."""
vtk_dir = os.path.join(source, "VTK")
search_root = vtk_dir if os.path.isdir(vtk_dir) else source
def step_index(path):
suffix = os.path.basename(os.path.dirname(path)).rsplit("_", 1)[-1]
return int(suffix) if suffix.isdigit() else 0
vtk_files = sorted(
glob.glob(os.path.join(search_root, "**", "internal.vtu"), recursive=True),
key=step_index
)
if not vtk_files:
raise FileNotFoundError(
f"No internal.vtu files found under: {search_root}. "
"Run foamToVTK first or pass a .vtu/.vtk file directly."
)
return vtk_files
def _resolve_openfoam_variables(requested, point_data_keys):
"""
Map requested variable names to OpenFOAM point-data fields.
Returns a list of (requested_name, field_name, component) with
component = None for scalars or vector magnitudes, or 0/1/2 for a
vector component.
"""
lower_map = {key.lower(): key for key in point_data_keys}
component_aliases = {
"velocity u": 0, "u_x": 0, "ux": 0,
"velocity v": 1, "u_y": 1, "uy": 1,
"velocity w": 2, "u_z": 2, "uz": 2,
}
resolved = []
for name in requested:
key = str(name).strip().lower()
if key == "velocity":
if "u" not in lower_map:
raise ValueError(
"Cannot resolve 'velocity': field 'U' not found. "
f"Available fields: {sorted(point_data_keys)}"
)
resolved.append((name, lower_map["u"], None))
elif key in component_aliases:
if "u" not in lower_map:
raise ValueError(
f"Cannot resolve '{name}': field 'U' not found. "
f"Available fields: {sorted(point_data_keys)}"
)
resolved.append((name, lower_map["u"], component_aliases[key]))
elif key in lower_map:
resolved.append((name, lower_map[key], None))
else:
raise ValueError(
f"Field '{name}' not found in the VTK output. "
f"Available fields: {sorted(point_data_keys)}"
)
return resolved
def _extract_openfoam(source, variables, x_pts, y_pts, z_pts,
time, time_index, n_last, method, k):
import pyvista as pv
if z_pts is None:
raise ValueError(
"OpenFOAM results are three-dimensional: provide the z "
"coordinate(s) of the extraction point(s)."
)
if os.path.isdir(source):
vtk_files = _collect_vtk_files(source)
if time == "last":
selected_files = vtk_files[-1:]
elif time == "index":
selected_files = [vtk_files[time_index]]
else:
selected_files = vtk_files[-n_last:]
else:
if not os.path.isfile(source):
raise FileNotFoundError(f"VTK file not found: {source}")
selected_files = [source]
mesh0 = pv.read(selected_files[0])
coords = np.asarray(mesh0.points)
resolved_variables = _resolve_openfoam_variables(
variables, list(mesh0.point_data.keys())
)
field_names = sorted({field for _, field, _ in resolved_variables})
# Average the requested fields over the selected time steps.
field_data = {
field: np.asarray(mesh0.point_data[field], dtype=float).copy()
for field in field_names
}
for vtk_file in selected_files[1:]:
mesh = pv.read(vtk_file)
for field in field_names:
field_data[field] += np.asarray(mesh.point_data[field], dtype=float)
for field in field_names:
field_data[field] /= len(selected_files)
tree = spatial.cKDTree(coords)
k_use = 1 if method == "nearest" else min(k, len(coords))
records = []
for i in range(x_pts.size):
target = np.array([x_pts[i], y_pts[i], z_pts[i]])
_, idx = tree.query(target, k=k_use)
idx = np.atleast_1d(idx)
weights = _idw_weights(coords[idx], target) if k_use > 1 else np.array([1.0])
record = {"point": i, "x": x_pts[i], "y": y_pts[i], "z": z_pts[i]}
for name, field, component in resolved_variables:
values = field_data[field][idx]
if values.ndim == 1:
record[name] = float(weights @ values)
elif component is not None:
record[name] = float(weights @ values[:, component])
else:
# Vector field requested without component: magnitude.
record[name] = float(
np.linalg.norm(weights @ values)
)
records.append(record)
return pd.DataFrame(records)