Skip to content

Unit-aware labels in visualisation functions #118

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lmoresi merged 2 commits into from
Apr 28, 2026
Merged

Unit-aware labels in visualisation functions #118

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3d8d62b
Add unit-aware labels to visualisation functions
lmoresi Apr 13, 2026
a62c18c
Document arrow scaling issue with physical units in plot_vector
lmoresi Apr 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 52 additions & 39 deletions src/underworld3/visualisation/visualisation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -343,7 +343,7 @@ def meshVariable_to_pv_mesh_object(meshVar, alpha=None):
return pv_mesh


def scalar_fn_to_pv_points(pv_mesh, uw_fn, dim=None, simplify=True):
def scalar_fn_to_pv_points(pv_mesh, uw_fn, dim=None):

Copilot AI Apr 17, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

scalar_fn_to_pv_points is part of the public underworld3.visualisation API (re-exported in src/underworld3/visualisation/__init__.py). Removing the simplify parameter is a breaking change for any downstream code calling it with simplify= or a 4th positional arg. Consider keeping the parameter for backward compatibility (e.g., simplify: bool = False), and either ignore it or deprecate it while avoiding sympy.simplify() by default.

Copilot uses AI. Check for mistakes.
"""Evaluate Underworld scalar function at PyVista mesh points.

Parameters
Expand All @@ -354,51 +354,49 @@ def scalar_fn_to_pv_points(pv_mesh, uw_fn, dim=None, simplify=True):
Underworld scalar function to evaluate.
dim : int, optional
Dimensionality (2 or 3). Auto-detected if None.
simplify : bool, optional
Simplify expression before evaluation (default: True).

Returns
-------
numpy.ndarray
Scalar values at mesh points.
Scalar values at mesh points (units stripped for PyVista).
The units string is stored as ``pv_mesh._last_scalar_units``.
"""
import underworld3 as uw
import sympy
import numpy as np

if simplify:
uw_fn = sympy.simplify(uw_fn)

if dim is None:
if pv_mesh.points[:, 2].max() - pv_mesh.points[:, 2].min() < 1.0e-6:
dim = 2
else:
dim = 3

# Use stored coordinate array if available (preserves units and dimensional info)
# Otherwise fall back to pv_mesh.points (non-dimensional [0-1] coordinates)
if hasattr(pv_mesh, '_coord_array'):
# Use the original mesh coordinate array (may be UnitAwareArray)
coords = pv_mesh._coord_array[:, 0:dim]
else:
# Fallback: use PyVista points directly (non-dimensional)
coords = pv_mesh.points[:, 0:dim]

scalar_values = uw.function.evaluate(uw_fn, coords, evalf=True)
scalar_values = uw.function.evaluate(uw_fn, coords)

# Capture units before stripping for colorbar labels
scalar_units = None
if hasattr(scalar_values, "units") and scalar_values.units is not None:
scalar_units = str(scalar_values.units)
elif hasattr(scalar_values, "_units") and scalar_values._units is not None:
scalar_units = str(scalar_values._units)

# Convert UnitAwareArray to plain numpy array for PyVista compatibility
# PyVista doesn't support UnitAwareArray and calls np.ndim() which fails
pv_mesh._last_scalar_units = scalar_units

# Strip units for PyVista compatibility
if hasattr(scalar_values, "magnitude"):
# UnitAwareArray - strip units for visualization
scalar_values = scalar_values.magnitude
else:
# Plain array - ensure it's numpy
scalar_values = np.asarray(scalar_values)

return scalar_values


def vector_fn_to_pv_points(pv_mesh, uw_fn, dim=None, simplify=True):
def vector_fn_to_pv_points(pv_mesh, uw_fn, dim=None):

Copilot AI Apr 17, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

vector_fn_to_pv_points is re-exported as part of the public underworld3.visualisation API. Removing the simplify parameter is an API break for external callers. To preserve compatibility, consider keeping simplify in the signature (defaulting to False) and treating it as a no-op / deprecated option rather than removing it outright.

Copilot uses AI. Check for mistakes.
"""Evaluate Underworld vector function at PyVista mesh points.

Parameters
Expand All @@ -409,41 +407,39 @@ def vector_fn_to_pv_points(pv_mesh, uw_fn, dim=None, simplify=True):
Underworld vector function to evaluate.
dim : int, optional
Dimensionality (not used, derived from function shape).
simplify : bool, optional
Simplify expression before evaluation (default: True).

Returns
-------
numpy.ndarray
Vector values at mesh points, shape ``(n_points, 3)``.
Units string stored as ``pv_mesh._last_vector_units``.
"""
import numpy as np
import underworld3 as uw
import sympy

if simplify:
uw_fn = sympy.simplify(uw_fn)
dim = uw_fn.shape[1]
if dim != 2 and dim != 3:
print(f"UW vector function should have dimension 2 or 3")

# Use stored coordinate array if available (preserves units and dimensional info)
# Otherwise fall back to pv_mesh.points (non-dimensional [0-1] coordinates)
if hasattr(pv_mesh, '_coord_array'):
# Use the original mesh coordinate array (may be UnitAwareArray)
coords = pv_mesh._coord_array[:, 0:dim]
else:
# Fallback: use PyVista points directly (non-dimensional)
coords = pv_mesh.points[:, 0:dim]

vector_values_raw = uw.function.evaluate(uw_fn, coords, evalf=True)
vector_values_raw = uw.function.evaluate(uw_fn, coords)

# Capture units before stripping
vector_units = None
if hasattr(vector_values_raw, "units") and vector_values_raw.units is not None:
vector_units = str(vector_values_raw.units)
elif hasattr(vector_values_raw, "_units") and vector_values_raw._units is not None:
vector_units = str(vector_values_raw._units)

pv_mesh._last_vector_units = vector_units

# Convert UnitAwareArray to plain numpy array for PyVista compatibility
if hasattr(vector_values_raw, "magnitude"):
# UnitAwareArray - strip units for visualization
vector_values_raw = vector_values_raw.magnitude
else:
# Plain array - ensure it's numpy
vector_values_raw = np.asarray(vector_values_raw)

vector_values = np.zeros_like(pv_mesh.points)
Expand Down Expand Up @@ -670,7 +666,10 @@ def plot_scalar(
pvmesh = mesh_to_pv_mesh(mesh)
pvmesh.point_data[scalar_name] = scalar_fn_to_pv_points(pvmesh, scalar)

print(pvmesh.point_data[scalar_name].min(), pvmesh.point_data[scalar_name].max())
# Build scalar bar label with units if available
scalar_bar_title = scalar_name
if hasattr(pvmesh, '_last_scalar_units') and pvmesh._last_scalar_units:
scalar_bar_title = f"{scalar_name} ({pvmesh._last_scalar_units})"

pl = pv.Plotter(window_size=window_size)
if clip_angle != 0.0:
Expand All @@ -683,7 +682,8 @@ def plot_scalar(
scalars=scalar_name,
show_edges=show_edges,
use_transparency=False,
show_scalar_bar=False,
show_scalar_bar=True,
scalar_bar_args={"title": scalar_bar_title},
opacity=1.0,
clim=clim,
)
Expand All @@ -697,7 +697,8 @@ def plot_scalar(
use_transparency=False,
opacity=1.0,
clim=clim,
show_scalar_bar=False,
show_scalar_bar=True,
scalar_bar_args={"title": scalar_bar_title},
)
Comment on lines 698 to 702

Copilot AI Apr 17, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

plot_scalar now forces show_scalar_bar=True (in both clipped and unclipped paths). This is a visible behavior change compared to the previous default and appears to contradict the PR description claim that unitless scripts render unchanged. If the intent is only to add unit-aware titles, consider keeping the prior scalar-bar visibility behavior (or add a show_scalar_bar parameter defaulting to the old value), and only pass scalar_bar_args when the scalar bar is actually shown. Also note scalar_name defaults to "", so enabling the scalar bar can result in an empty title in existing calls that didn't pass a name.

Copilot uses AI. Check for mistakes.

pl.show(cpos=cpos)
Expand Down Expand Up @@ -811,6 +812,16 @@ def plot_vector(
None
This function does not return any value. It displays the vector field on the mesh in a PyVista
window and optionally saves a screenshot.

Notes
-----
When the model uses physical units, arrows are drawn in the mesh
coordinate space. A velocity of 1 cm/yr on a mesh in meters
(extent ~1e6) produces arrows of length ~3e-10 in mesh units —
effectively invisible. Adjust ``vmag`` to compensate::

# Scale arrows to ~5% of mesh extent
vmag = 0.05 * mesh_extent / max_velocity
"""

import sympy
Expand All @@ -827,7 +838,10 @@ def plot_vector(
else:
pvmesh.point_data[scalar_name] = scalar_fn_to_pv_points(pvmesh, scalar.sym)

print(pvmesh.point_data[scalar_name].min(), pvmesh.point_data[scalar_name].max())
# Build scalar bar label with units if available
scalar_bar_title = scalar_name
if hasattr(pvmesh, '_last_scalar_units') and pvmesh._last_scalar_units:
scalar_bar_title = f"{scalar_name} ({pvmesh._last_scalar_units})"

velocity_points = meshVariable_to_pv_cloud(vector)
velocity_points.point_data[vector_name] = vector_fn_to_pv_points(velocity_points, vector.sym)
Expand All @@ -843,7 +857,8 @@ def plot_vector(
scalars=scalar_name,
show_edges=show_edges,
use_transparency=False,
show_scalar_bar=False,
show_scalar_bar=True,
scalar_bar_args={"title": scalar_bar_title},
opacity=1.0,
clim=clim,
)
Expand All @@ -857,12 +872,10 @@ def plot_vector(
use_transparency=False,
opacity=1.0,
clim=clim,
show_scalar_bar=False,
show_scalar_bar=True,
scalar_bar_args={"title": scalar_bar_title},
)
Comment on lines 873 to 877

Copilot AI Apr 17, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

plot_vector now forces show_scalar_bar=True (clipped and unclipped paths). This changes the default rendered output for unitless runs (the PR description says those should remain unchanged). If the goal is to improve the scalar-bar title when units are active, consider restoring the prior default scalar-bar visibility (or add an explicit parameter), and only supply scalar_bar_args when the scalar bar is being shown.

Copilot uses AI. Check for mistakes.

# pl.add_scalar_bar(vector_name, vertical=False, title_font_size=25, label_font_size=20, fmt=fmt,
# position_x=0.225, position_y=0.01,)

if show_arrows:
pl.add_arrows(
velocity_points.points[::vfreq],
Expand Down
Loading