"""
Underworld3: Finite element geodynamics with symbolic expressions.
Underworld3 is a finite element, particle-in-cell geodynamics code that
produces mathematically self-describing models through an interface with
SymPy. It builds upon PETSc for parallel finite element assembly and
solving, using petsc4py for Python integration.
A common pattern for building Underworld3 models is to develop Python
scripts in notebook-friendly form (e.g., with jupytext) which are
documented through markdown descriptions. Underworld objects are similarly
documented so that their underlying algorithmic and mathematical structure
can be examined in a notebook.
Scripts built this way are also compatible with ``mpirun`` for parallel
execution.
Documentation
-------------
The Underworld documentation has two parts:
- **User/Theory Manual**: https://underworldcode.github.io/underworld3/main/
- **API Reference**: https://underworldcode.github.io/underworld3/main_api/
The development branch has similar documentation at ``/development/``.
Installation
------------
From the repository::
./scripts/compile.sh
For development (editable install)::
pip install -e . --no-build-isolation
To uninstall::
pip uninstall underworld3
Testing
-------
Run the pytest suite::
./scripts/test.sh
Build API documentation::
./scripts/build_api_docs.sh
Then open ``uw3_api_docs/index.html`` to browse.
Key Submodules
--------------
meshing
Mesh generation functions (boxes, shells, annuli).
discretisation
Mesh and MeshVariable classes.
systems
PDE solvers (Stokes, Poisson, advection-diffusion).
constitutive_models
Material rheology definitions.
function
Symbolic expression evaluation.
swarm
Particle swarm management.
maths
Vector calculus and tensor operations.
"""
from mpi4py import MPI # for initialising MPI
import petsc4py as _petsc4py
import sys
import os as _os
import warnings as _warnings
_petsc4py.init(sys.argv)
from petsc4py import PETSc
# pop the default petsc Signal handler to let petsc errors appear in python
# unclear if this is the appropriate way see discussion
# https://gitlab.com/petsc/petsc/-/issues/1066
try:
PETSc.Sys.popErrorHandler()
options = PETSc.Options()
options["options_left"] = 0
except (TypeError, AttributeError):
# Handle Sphinx autodoc mocking - PETSc mock objects don't support these operations
pass
def _parse_thread_env(name: str):
"""Return integer thread count from environment variable, or None."""
raw = _os.environ.get(name)
if raw is None or str(raw).strip() == "":
return None
try:
return int(float(raw))
except (TypeError, ValueError):
return None
def _apply_mpi_thread_policy():
"""
Default thread policy for MPI runs.
Caps common thread pools to 1 unless explicitly set by the user. This
prevents MPI+BLAS oversubscription (large performance regressions).
"""
if MPI.COMM_WORLD.Get_size() <= 1:
return
if _os.environ.get("UW_DISABLE_THREAD_CAPS", "0").lower() in ("1", "true", "yes", "on"):
return
thread_vars = (
"OMP_NUM_THREADS",
"OPENBLAS_NUM_THREADS",
"MKL_NUM_THREADS",
"VECLIB_MAXIMUM_THREADS",
"NUMEXPR_NUM_THREADS",
)
for var in thread_vars:
if _os.environ.get(var, "").strip() == "":
_os.environ[var] = "1"
if MPI.COMM_WORLD.Get_rank() == 0 and _os.environ.get("UW_SUPPRESS_THREAD_WARNING", "0").lower() not in (
"1",
"true",
"yes",
"on",
):
oversub = {var: _parse_thread_env(var) for var in thread_vars}
oversub = {k: v for k, v in oversub.items() if v is not None and v > 1}
if oversub:
items = ", ".join(f"{k}={v}" for k, v in oversub.items())
_warnings.warn(
"MPI run with thread pools > 1 detected "
f"({items}). This may cause severe oversubscription. "
"Set thread counts to 1 per rank, or set UW_SUPPRESS_THREAD_WARNING=1 "
"if this is intentional.",
RuntimeWarning,
stacklevel=2,
)
_apply_mpi_thread_policy()
# Version is derived from git tags via setuptools_scm
# Priority: 1) _version.py (generated by setuptools_scm at build time)
# 2) importlib.metadata (for installed packages)
# 3) Fallback for development environments
try:
from ._version import __version__
except ImportError:
try:
from importlib.metadata import version as _get_version
__version__ = _get_version("underworld3")
except Exception:
__version__ = "0.0.0+unknown" # Fallback for non-git, non-installed environments
def view():
from IPython.display import Latex, Markdown, display
from textwrap import dedent
import inspect
## Docstring (static)
docstring = dedent(__doc__)
# docstring = docstring.replace("$", "$").replace("$", "$")
display(Markdown(docstring))
return
# PETSc.Log().begin()
# Bundle these utils
# Needed everywhere
import underworld3.mpi
from .mpi import pprint, selective_ranks, collective_operation, CollectiveOperationError
from ._var_types import *
from .utilities._petsc_tools import *
from .utilities._nb_tools import *
from .utilities._utils import auditor
from .utilities import _api_tools
# from underworld3.utilities import _api_tools
from .utilities._utils import auditor
import underworld3.adaptivity
import underworld3.coordinates
import underworld3.discretisation
import underworld3.meshing
import underworld3.constitutive_models
import underworld3.function
import underworld3.maths
import underworld3.swarm
import underworld3.systems
import underworld3.maths
import underworld3.utilities
import underworld3.model
import underworld3.parameters
import underworld3.materials
import underworld3.discretisation.persistence
import underworld3.checkpoint
from .model import (
Model,
create_model,
get_default_model,
reset_default_model,
ThermalConvectionConfig,
create_thermal_convection_model,
)
from .parameters import ParameterRegistry, ParameterType
from .materials import MaterialRegistry, MaterialProperty
from .constitutive_models import MultiMaterialConstitutiveModel
from .function import quantity, expression, with_units, expand, unwrap
from .coordinates import uwdiff # Differentiation helper for UWCoordinates
from .utilities import retention_curves
# Unit utilities (top-level convenience for user code)
from .function.unit_conversion import _extract_value
# Currently on binder, pykdtree is hanging - fallback to previous implementation
# import underworld3.kdtree
import underworld3.ckdtree as kdtree
import underworld3.cython
import underworld3.scaling
import underworld3.visualisation
import numpy as _np
# REMOVED: Deprecated Units mixin system (2025-11-01)
# The UnitAwareMixin system was never integrated and is now deprecated
# Active implementation: discretisation/enhanced_variables.py
# Historical code preserved in: utilities/units_mixin.py (not exported)
# High-level units utilities
from . import units as _units_module # Keep module for internal use
from .units import (
check_units_consistency,
get_dimensionality,
assert_dimensionality, # General type-safety gatekeeper for dimensionality
units_of,
get_units,
non_dimensionalise,
dimensionalise,
create_quantity,
convert_units,
to_base_units,
to_reduced_units,
to_compact,
is_dimensionless,
has_units,
same_units,
validate_expression_units,
enforce_units_consistency,
# derivative_units removed - use natural pint arithmetic: var.units / mesh.units
)
# Universal unit query function is already imported from units module (line 172)
# NOTE: get_units has been moved to units module for unified API
# Make units registry easily accessible (remove scaling terminology)
# Users can now do uw.units.K, uw.units.Pa, etc. instead of uw.scaling.units.K
from .scaling import units
# Add view() method to units registry following established pattern
def _units_view():
"""Display units registry information following the established view() pattern."""
try:
from IPython.display import Markdown, display
content = f"""## Units Registry
The Underworld3 units registry provides access to physical units for dimensional analysis.
### Common Units Examples:
- **Temperature**: `uw.units.K`, `uw.units.celsius`, `uw.units.degC`
- **Pressure**: `uw.units.Pa`, `uw.units.bar`, `uw.units.atm`
- **Length**: `uw.units.m`, `uw.units.cm`, `uw.units.km`
- **Time**: `uw.units.s`, `uw.units.year`, `uw.units.Ma` (million years)
- **Viscosity**: `uw.units.Pa * uw.units.s`
- **Velocity**: `uw.units.cm / uw.units.year`
### Usage:
```python
# Create quantities
temperature = 1500 * uw.units.K
viscosity = 1e21 * uw.units.Pa * uw.units.s
velocity = 5 * uw.units.cm / uw.units.year
# Set model reference quantities
model.set_reference_quantities(
mantle_temperature=temperature,
mantle_viscosity=viscosity
)
```
**Total units available**: {len([attr for attr in dir(units) if not attr.startswith('_')])}
"""
display(Markdown(content))
except ImportError:
# Fallback for non-Jupyter environments
print("Underworld3 Units Registry")
print("=" * 30)
print(
f"Total units available: {len([attr for attr in dir(units) if not attr.startswith('_')])}"
)
print("\nCommon units:")
print(" Temperature: uw.units.K, uw.units.celsius")
print(" Pressure: uw.units.Pa, uw.units.bar")
print(" Length: uw.units.m, uw.units.cm, uw.units.km")
print(" Time: uw.units.s, uw.units.year")
print("\nUsage:")
print(" temperature = 1500 * uw.units.K")
print(" viscosity = 1e21 * uw.units.Pa * uw.units.s")
# Attach view method to units registry
units.view = _units_view
# ============================================================================
# NON-DIMENSIONAL SCALING - SINGLE GLOBAL FLAG
# ============================================================================
# This is the ONLY place where non-dimensional scaling is enabled/disabled.
# All code checks this flag via is_nondimensional_scaling_active().
# ============================================================================
from contextlib import contextmanager
# SINGLE GLOBAL FLAG for non-dimensional scaling
_USE_NONDIMENSIONAL_SCALING = False
def use_nondimensional_scaling(enabled=True):
"""
Enable or disable non-dimensional scaling globally.
When enabled, equations are scaled during unwrap() for better numerical
conditioning. Variables with units and reference scales are automatically
scaled (divided by their scaling_coefficient).
This is the ONLY way to control non-dimensional scaling - there are no
other flags, solver settings, or context managers that affect this.
Parameters
----------
enabled : bool, default=True
True to enable non-dimensional scaling, False to disable
Notes
-----
- Scaling coefficients are ALWAYS computed from model.set_reference_quantities()
- This flag only controls whether those coefficients are APPLIED during unwrap()
- Changing this flag requires recompiling solvers (set solver.is_setup=False)
Examples
--------
Setup problem with reference quantities:
>>> model = uw.Model()
>>> model.set_reference_quantities(
... domain_depth=uw.quantity(100, "km"),
... temperature_diff=uw.quantity(1000, "kelvin")
... )
>>> T = uw.discretisation.MeshVariable('T', mesh, 1, units='kelvin')
Solve with dimensional form (default):
>>> uw.use_nondimensional_scaling(False) # Default
>>> poisson.solve()
>>> u_dimensional = T.array.copy()
Solve with non-dimensional form (for comparison):
>>> uw.use_nondimensional_scaling(True)
>>> poisson.is_setup = False # Force recompilation with new scaling
>>> poisson.solve()
>>> u_nondimensional = T.array.copy()
>>>
>>> # Solutions should be identical
>>> import numpy as np
>>> assert np.allclose(u_dimensional, u_nondimensional)
See Also
--------
is_nondimensional_scaling_active : Check current scaling state
model.set_reference_quantities : Set reference scales
"""
global _USE_NONDIMENSIONAL_SCALING
_USE_NONDIMENSIONAL_SCALING = bool(enabled)
[docs]
def is_nondimensional_scaling_active():
"""
Check if non-dimensional scaling is currently enabled.
Returns
-------
bool
True if non-dimensional scaling is active, False otherwise
See Also
--------
use_nondimensional_scaling : Enable/disable scaling
"""
return _USE_NONDIMENSIONAL_SCALING
@contextmanager
def nondimensional_scaling_context(enabled=True):
"""
Context manager to temporarily set non-dimensional scaling state.
Automatically restores the previous state when exiting the context.
This is useful for operations that need to temporarily work in
non-dimensional space (like semi-Lagrangian advection) without
affecting the global state.
Parameters
----------
enabled : bool, default=True
True to enable non-dimensional scaling within context
Examples
--------
>>> # Temporarily enable non-dimensional mode
>>> with uw.nondimensional_scaling_context(True):
... # evaluate() returns non-dimensional results here
... velocity_nd = uw.function.evaluate(u.sym, coords)
>>> # Original state restored automatically
See Also
--------
use_nondimensional_scaling : Permanently set scaling state
is_nondimensional_scaling_active : Check current state
"""
global _USE_NONDIMENSIONAL_SCALING
old_value = _USE_NONDIMENSIONAL_SCALING
_USE_NONDIMENSIONAL_SCALING = bool(enabled)
try:
yield
finally:
_USE_NONDIMENSIONAL_SCALING = old_value
# ============================================================================
# STRICT UNITS MODE - ENFORCES UNITS-SCALES CONTRACT
# ============================================================================
# This flag enforces the principle: "Units require reference quantities"
# When enabled, variables with units REQUIRE reference quantities to be set.
# DEFAULT: ON - enforces best practices from the start
# ============================================================================
# SINGLE GLOBAL FLAG for strict units enforcement
_STRICT_UNITS_MODE = True # Default: ON (enforces units-scales contract)
def use_strict_units(enabled=True):
"""
Enable or disable strict units enforcement.
**DEFAULT: ON** - Strict units mode is enabled by default to enforce best
practices from the start. Variables with units REQUIRE reference quantities.
When disabled, variables with units are allowed without reference quantities
but will get scaling_coefficient=1.0 and a warning. This "half-way zone"
leads to poor numerical conditioning and silent errors.
**Disabling strict mode is only recommended for:**
- Expert users who understand the implications
- Debugging specific issues
- Legacy code migration (temporary)
Parameters
----------
enabled : bool, default=True
True to enforce strict units (DEFAULT - recommended)
False to allow units without reference quantities (expert/debugging only)
Examples
--------
Normal usage (strict mode ON by default):
>>> import underworld3 as uw
>>> # Strict mode is ON by default - no need to enable
>>>
>>> # Set reference quantities FIRST:
>>> model = uw.get_default_model()
>>> model.set_reference_quantities(
... domain_depth=uw.quantity(1000, 'km'),
... plate_velocity=uw.quantity(5, 'cm/year')
... )
>>>
>>> # Then create mesh and variables:
>>> mesh = uw.meshing.StructuredQuadBox(elementRes=(4, 4))
>>> v = uw.discretisation.MeshVariable("v", mesh, 2, units="m/s") # ✓ OK
Disable for debugging (expert use only):
>>> uw.use_strict_units(False) # Expert/debugging only
>>> mesh = uw.meshing.StructuredQuadBox(elementRes=(4, 4))
>>> v = uw.discretisation.MeshVariable("v", mesh, 2, units="m/s") # ⚠️ Warning
See Also
--------
is_strict_units_active : Check current strict mode
Model.set_reference_quantities : Set reference quantities
"""
global _STRICT_UNITS_MODE
_STRICT_UNITS_MODE = bool(enabled)
def is_strict_units_active():
"""
Check if strict units enforcement is enabled.
Returns
-------
bool
True if strict units mode is active, False otherwise
See Also
--------
use_strict_units : Enable/disable strict units
"""
return _STRICT_UNITS_MODE
# ============================================================================
# INTERNAL COMPATIBILITY - Do not use directly
# ============================================================================
# These exist for backward compatibility with old scaling experiments.
# New code should use use_nondimensional_scaling() instead.
@contextmanager
def _apply_scaling():
"""Internal context manager - DEPRECATED, use use_nondimensional_scaling() instead."""
global _USE_NONDIMENSIONAL_SCALING
old_value = _USE_NONDIMENSIONAL_SCALING
_USE_NONDIMENSIONAL_SCALING = True
try:
yield
finally:
_USE_NONDIMENSIONAL_SCALING = old_value
def _is_scaling_active():
"""Internal check - DEPRECATED, use is_nondimensional_scaling_active() instead."""
return _USE_NONDIMENSIONAL_SCALING
# Backward compatibility aliases (deprecated)
scaled_symbols = (
_apply_scaling # For backward compatibility - use use_nondimensional_scaling() instead
)
def unwrap(fn, keep_constants=True, return_self=True, apply_scaling=False):
"""
Unwrap expressions with optional automatic scaling.
Parameters:
-----------
fn : expression
The expression to unwrap
keep_constants : bool, default=True
Whether to keep constants in the unwrapped expression
return_self : bool, default=True
Whether to return self if no unwrapping is needed
apply_scaling : bool, default=False
Whether to automatically apply scale factors to variables with units
Example:
model = uw.Model()
model.set_reference_quantities(mantle_temperature=1500*uw.scaling.units.K)
temperature = uw.discretisation.MeshVariable("T", mesh, 1, units="K")
expr = uw.function.expression("heat", 2 * temperature.sym, "heat equation")
# Normal unwrap
result = uw.unwrap(expr)
# Unwrap with automatic scaling
scaled_result = uw.unwrap(expr, apply_scaling=True)
"""
from .function.expressions import unwrap as _unwrap
if apply_scaling:
with _apply_scaling():
return _unwrap(fn, keep_constants, return_self)
else:
return _unwrap(fn, keep_constants, return_self)
# Enhanced variables with units and mathematical operations (internal use only)
from .discretisation.enhanced_variables import (
# EnhancedMeshVariable, # Hidden - use uw.discretisation.MeshVariable
create_enhanced_mesh_variable,
# Note: SwarmVariable already has enhanced functionality built-in (see swarm.py)
# No separate EnhancedSwarmVariable needed - use uw.swarm.SwarmVariable directly
)
# Make MeshVariable available at top level (this is the enhanced version)
from .discretisation import MeshVariable
def synchronised_array_update(context_info="user operations"):
"""
Context manager for synchronised array updates across multiple variables.
Batches multiple array assignments together and defers PETSc synchronization
until the end of the context, ensuring atomic updates and better performance.
Example
-------
with uw.synchronised_array_update():
velocity.array[...] = new_velocity_values
pressure.array[...] = new_pressure_values
temperature.array[...] = new_temperature_values
# All arrays are synchronized here
Parameters
----------
context_info : str
Optional description of the update context for debugging
Returns
-------
Context manager for delayed callback execution
"""
return utilities.NDArray_With_Callback.delay_callbacks_global(context_info)
# get the id for this installation of uw, if not there, create it
try:
from ._uwid import uwid as _id
except:
import uuid as _uuid
_id = str(_uuid.uuid4())
###########################################
####### For User telemetry metrics ########
# setup metrics function *only* if we are rank==0.
# Metric can be disables by setting UW_NO_USAGE_METRICS env var.
# TODO: confirm not other condition like tests?
if underworld3.mpi.rank == 0:
def _sendData():
import os
# disable collection of data if requested
if "UW_NO_USAGE_METRICS" not in os.environ:
# get platform info
import platform
sysinfo = platform.system()
sysinfo += "__" + platform.machine()
# check if docker
if os.path.isfile("/.dockerinit"):
sysinfo += "__docker"
event_dict = {
"properties": {
"version": underworld3.__version__,
"platform": sysinfo,
"run_size": underworld3.mpi.size,
"distinct_id": _id,
}
}
# send info async
import threading
thread = threading.Thread(
target=underworld3.utilities._utils.postHog,
args=("import_uw3", event_dict),
)
thread.daemon = True
thread.start()
try:
_sendData()
except: # continue quietly if something above failed
pass
####### END User telemetry metrics ########
# Info for JIT modules.
# These dicts should be populated by submodules
# which define cython/c based classes.
# We use ordered dictionaries because the
# ordering can be important when linking in libraries.
# Note that actually what we want is an ordered set (which Python
# doesn't natively provide). Hence for the key/value pair,
# the value is always set to `None`.
from collections import OrderedDict as _OD
_libfiles = _OD()
_libdirs = _OD()
_incdirs = _OD({_np.get_include(): None})
## -------------------------------------------------------------
# DIAGNOSTICS - uw.doctor() for detecting configuration issues
## -------------------------------------------------------------
from .utilities.diagnostics import doctor, health_check
## -------------------------------------------------------------
# PARAMETERS - uw.Params() for notebook/CLI configurable parameters
# uw.Param() wrapper for typed parameters with units support
## -------------------------------------------------------------
from .utilities._params import Params, Param, ParamType
## -------------------------------------------------------------
# PAUSE - Interactive pause for notebook development
## -------------------------------------------------------------
from .utilities._interrupt import UW_Pause, pause, is_notebook
## -------------------------------------------------------------
# pdoc3 over-rides. pdoc3 has a strange path-traversal algorithm
# that seems to have trouble finding modules if we move this
# dictionary to any other location in the underworld3 tree
__pdoc__ = {}
# Cython files cannot be documented. We should move pure
# python out of these files if we can
__pdoc__["cython"] = False
# __pdoc__["function.analytic"] = False
# Here we nuke some of the placeholders in the parent class so that they do not mask the
# child class modifications
__pdoc__["systems.constitutive_models.Constitutive_Model.Parameters"] = False
# Note: SymPy converter registration approach doesn't work reliably in strict mode
# The better approach is to ensure UWexpression arithmetic operations return SymPy objects
# This is handled by the __rmul__, __radd__ etc. methods in the mathematical mixin