geoprior.models.subsidence.maths#

GeoPrior maths helpers (physics terms + scaling).

Functions

bounded_exp(raw, log_min, log_max, *[, eps, ...])

Exponentiate a raw parameter inside hard log-bounds.

compose_physics_fields(model, *, ...[, ...])

Compose physically meaningful fields \(K\), \(S_s\), and \(\tau\) from network "base" logits and coordinate corrections.

compute_bounds_residual(model, *, H_field[, ...])

Compute differentiable bound-violation residuals for the learned physics fields.

compute_consistency_prior(model, K_field, ...)

Compute the consolidation timescale consistency prior.

compute_consolidation_residual(model, ds_dt, ...)

Consolidation PDE residual (Voigt).

compute_consolidation_step_residual(*, ...)

Compute a one-step consolidation residual in SI space.

compute_gw_flow_residual(model, dh_dt, ...)

Groundwater flow PDE residual (NaN/Inf-safe, broadcast-safe).

compute_mv_prior(model[, Ss_field, logSs, ...])

Compute an m_v - gamma_w prior from predicted S_s.

compute_scales(model, *, t, s_mean, h_mean, ...)

Compute robust normalization scales for physics residuals.

compute_smoothness_prior(dK_dx, dK_dy, ...)

Compute a smoothness prior on spatial gradients of physics fields.

cons_step_to_cons_residual(cons_step_m, *, ...)

Convert consolidation step residual (meters per step) into the chosen residual units.

dt_to_seconds(dt, *, time_units)

dt(time_units) -> seconds.

ensure_3d(x)

Return a rank-3 tensor, preferring static rank when available.

equilibrium_compaction_si(*, h_mean_si, ...)

Compute equilibrium compaction s_eq in SI meters.

exp_from_bounds(raw_log, log_min, log_max, *)

finite_floor(x, eps)

Replace NaN/Inf by eps and floor to eps.

get_log_bounds(model, *[, as_tensor, dtype, ...])

Get validated log-space bounds for K and Ss.

get_log_tau_bounds(model, *[, as_tensor, ...])

Get validated log-space bounds for the consolidation timescale.

guard_scale_with_residual(residual, scale, ...)

Guard a residual scale using the observed residual magnitude.

huber(x, *[, delta])

Huber loss (elementwise).

integrate_consolidation_mean(*, h_mean_si, ...)

Integrate mean consolidation settlement over a forecast horizon.

normalize_time_units(u)

Normalize time unit strings.

positive(x, *[, eps])

Softplus positivity.

q_to_gw_source_term_si(model, Q_logits, *, ...)

Convert Q_logits into a GW source term in SI units.

q_to_per_second(Q_base, *, scaling_kwargs, ...)

Normalize Q into 1/s.

rate_to_per_second(dz_dt, *, time_units)

d/d(time_units) -> d/ds.

resolve_auto_scale_floor(key, scaling_kwargs)

Robustly determine a numerical stability floor for physics scales.

resolve_cons_drawdown_options(scaling_kwargs, *)

Resolve consolidation drawdown options from scaling_kwargs.

resolve_cons_units(sk)

Normalize consolidation residual units.

resolve_gw_units(sk)

resolve_mv_gamma_log_target(model, Ss_field, *)

Return log(mv * gamma_w) with configurable units.

resolve_mv_gamma_log_target_from_logSs(...)

Like resolve_mv_gamma_log_target(), but uses logSs.

resolve_q_kind(sk)

Normalize Q meaning for gw forcing.

safe_log_pos(x, *[, eps, dtype])

log(safe_pos(x)).

safe_pos(x, *[, eps, dtype])

Force x to be finite and >= eps.

scale_residual(residual, scale, *[, floor])

Scale a residual by a (guarded) normalization factor.

seconds_per_time_unit(time_units, *[, dtype])

Seconds-per-unit.

settlement_state_for_pde(s_pred_si, t, *[, ...])

Map predicted settlement output into a PDE-ready settlement state.

smooth_relu(x, *[, beta])

Smooth approximation to relu(x) with controlled curvature.

tau_phys_from_fields(model, K_field, ...[, ...])

Compute the physics closure consolidation timescale tau_phys and the effective drainage thickness Hd.

tf_print_nonfinite(tag, x[, summarize])

Print a compact report ONLY if x contains NaN/Inf (graph-safe).

to_rms(x, *[, axis, keepdims, eps, ...])

Compute the root-mean-square (RMS) of a tensor.

vprint(verbose, *args)

Verbose print (eager-friendly).

Classes

LogClipConstraint(min_value, max_value)

NaN-safe clip constraint for log-parameters.

class geoprior.models.subsidence.maths.LogClipConstraint(min_value, max_value)[source]#

Bases: Constraint

NaN-safe clip constraint for log-parameters.

This constraint is intended for parameters stored in log-space, such as logK, logSs, or log_tau, where the model must enforce hard bounds:

(1)#\[w \in [w_{min}, w_{max}]\]
__init__(min_value, max_value)[source]#
geoprior.models.subsidence.maths.vprint(verbose, *args)[source]#

Verbose print (eager-friendly).

Parameters:

verbose (int)

Return type:

None

geoprior.models.subsidence.maths.tf_print_nonfinite(tag, x, summarize=6)[source]#

Print a compact report ONLY if x contains NaN/Inf (graph-safe).

Parameters:
  • tag (str)

  • x (Tensor)

  • summarize (int)

Return type:

Tensor

geoprior.models.subsidence.maths.resolve_q_kind(sk)[source]#

Normalize Q meaning for gw forcing.

Parameters:

sk (dict[str, Any] | None)

Return type:

str

geoprior.models.subsidence.maths.q_to_gw_source_term_si(model, Q_logits, *, Ss_field, H_field, coords_normalized, t_range_units, time_units, scaling_kwargs, H_floor=1.0, verbose=0)[source]#

Convert Q_logits into a GW source term in SI units.

This helper maps the network output Q_logits into a source term \(Q_{term}\) that is compatible with the groundwater PDE residual used by the model:

(2)#\[R_{gw} = S_s \, \frac{\partial h}{\partial t} - \nabla \cdot (K \nabla h) - Q_{term}\]

The returned tensor always has units of 1/s so it can be subtracted directly in \(R_{gw}\).

Parameters:
  • Q_logits (Tensor)

  • Ss_field (Any | None)

  • H_field (Any | None)

  • coords_normalized (bool)

  • t_range_units (Any | None)

  • time_units (str | None)

  • scaling_kwargs (dict[str, Any] | None)

  • H_floor (float)

  • verbose (int)

Return type:

Tensor

geoprior.models.subsidence.maths.q_to_per_second(Q_base, *, scaling_kwargs, time_units, coords_normalized, t_range_units=None, eps=1e-12)[source]#

Normalize Q into 1/s.

Assumed meaning (recommended default):
Q_kind = “per_volume” -> Q is already 1/time_unit or 1/s, representing

volumetric source/sink per unit volume.

If coords_normalized and Q_wrt_normalized_time=True, we de-normalize by the time range first (same chain rule as dh/dt).

Parameters:
  • Q_base (Tensor)

  • scaling_kwargs (dict[str, Any] | None)

  • time_units (str | None)

  • coords_normalized (bool)

  • t_range_units (Any | None)

  • eps (float)

Return type:

Tensor

geoprior.models.subsidence.maths.cons_step_to_cons_residual(cons_step_m, *, dt_units, scaling_kwargs, time_units, eps=1e-12)[source]#

Convert consolidation step residual (meters per step) into the chosen residual units. Supported outputs are "step" for meters, "time_unit" for meters per time unit, and "second" for meters per second (SI rate).

Parameters:
  • cons_step_m (Tensor)

  • dt_units (Tensor)

  • scaling_kwargs (dict[str, Any] | None)

  • time_units (str | None)

  • eps (float)

Return type:

Tensor

geoprior.models.subsidence.maths.resolve_mv_gamma_log_target_from_logSs(model, logSs, *, eps=1e-15, verbose=0)[source]#

Like resolve_mv_gamma_log_target(), but uses logSs.

This is the preferred path for mode=’logss’ because it avoids the 1/Ss gradient amplification from log(Ss_field).

Return type:

Tensor

geoprior.models.subsidence.maths.compute_mv_prior(model, Ss_field=None, *, logSs=None, mode=None, as_loss=True, weight=None, warmup_steps=None, step=None, alpha_disp=0.1, delta=1.0, eps=1e-15, verbose=0)[source]#

Compute an m_v - gamma_w prior from predicted S_s.

This routine builds a log-space residual that ties the model’s specific storage \(S_s\) to the consolidation coefficient \(m_v\) and the unit weight of water \(gamma_w\) via:

(3)#\[S_s \approx m_v \, \gamma_w\]

The constraint is applied in log space for numerical stability:

(4)#\[r = \log(S_s) - \log(m_v \, \gamma_w)\]

Depending on mode, gradients may be blocked or allowed to flow through \(S_s\) (or its log) to control stability.

Parameters:
  • Ss_field (Any | None)

  • logSs (Any | None)

  • mode (str | None)

  • as_loss (bool)

geoprior.models.subsidence.maths.resolve_mv_gamma_log_target(model, Ss_field, *, eps=1e-15, verbose=0)[source]#

Return log(mv * gamma_w) with configurable units.

If mv_prior_units == “strict”:

log_target = log(mv) + log(gamma_w)

If mv_prior_units == “auto”:

pick among 4 candidates that best matches mean(log(Ss_field)) in magnitude: - mv vs mv/1000 - gamma_w vs gamma_w*1000

Return type:

Tensor

geoprior.models.subsidence.maths.safe_pos(x, *, eps=1e-15, dtype=tf.float32)[source]#

Force x to be finite and >= eps.

Replaces NaN/Inf by eps, then floors.

geoprior.models.subsidence.maths.safe_log_pos(x, *, eps=1e-15, dtype=tf.float32)[source]#

log(safe_pos(x)).

geoprior.models.subsidence.maths.huber(x, *, delta=1.0)[source]#

Huber loss (elementwise).

delta is treated as a scalar constant.

geoprior.models.subsidence.maths.compute_gw_flow_residual(model, dh_dt, d_K_dh_dx_dx, d_K_dh_dy_dy, Ss_field, *, Q=None, verbose=0)[source]#

Groundwater flow PDE residual (NaN/Inf-safe, broadcast-safe).

Parameters:
  • dh_dt (Tensor)

  • d_K_dh_dx_dx (Tensor)

  • d_K_dh_dy_dy (Tensor)

  • Ss_field (Tensor)

  • Q (Any | None)

  • verbose (int)

Return type:

Tensor

geoprior.models.subsidence.maths.compute_consolidation_residual(model, ds_dt, s_state, h_mean, H_field, tau_field, *, Ss_field=None, inputs=None, verbose=0)[source]#

Consolidation PDE residual (Voigt).

Parameters:
  • ds_dt (Tensor)

  • s_state (Tensor)

  • h_mean (Tensor)

  • H_field (Tensor)

  • tau_field (Tensor)

  • Ss_field (Any | None)

  • inputs (dict[str, Tensor] | None)

  • verbose (int)

Return type:

Tensor

geoprior.models.subsidence.maths.equilibrium_compaction_si(*, h_mean_si, h_ref_si, Ss_field, H_field_si, drawdown_mode='smooth_relu', drawdown_rule='ref_minus_mean', relu_beta=20.0, stop_grad_ref=True, drawdown_zero_at_origin=False, drawdown_clip_max=None, eps=1e-15, verbose=0)[source]#

Compute equilibrium compaction s_eq in SI meters.

This function computes the equilibrium (instantaneous) settlement that would be reached under a sustained head change, given a specific storage field and a compressible thickness. The output s_eq is used by the consolidation residual to compare the current settlement state against its equilibrium target.

Parameters:
  • h_mean_si (Tensor)

  • h_ref_si (Tensor)

  • Ss_field (Tensor)

  • H_field_si (Tensor)

  • drawdown_mode (str)

  • drawdown_rule (str)

  • relu_beta (float)

  • stop_grad_ref (bool)

  • drawdown_zero_at_origin (bool)

  • drawdown_clip_max (float | None)

  • eps (float)

  • verbose (int)

Return type:

Tensor

geoprior.models.subsidence.maths.integrate_consolidation_mean(*, h_mean_si, Ss_field, H_field_si, tau_field, h_ref_si, s_init_si, dt=None, time_units='yr', method='exact', eps_tau=1e-12, relu_beta=20.0, drawdown_mode='smooth_relu', drawdown_rule='ref_minus_mean', stop_grad_ref=True, drawdown_zero_at_origin=False, drawdown_clip_max=None, verbose=0)[source]#

Integrate mean consolidation settlement over a forecast horizon.

This routine evolves the mean settlement state \(\bar{s}(t)\) using a stable, shape-safe time stepper that is compatible with TensorFlow graph execution. It is designed for the GeoPriorSubsNet “Option-1” mean path, where the mean subsidence is physics-driven from the predicted head.

The integrator advances a first-order relaxation model:

(5)#\[\frac{d\bar{s}}{dt} = \frac{s_{eq}(t) - \bar{s}(t)}{\tau(t)}\]

where:

  • \(\bar{s}(t)\) is the mean settlement state (m),

  • \(s_{eq}(t)\) is the equilibrium compaction (m),

  • \(\tau(t)\) is a consolidation time scale (s).

The equilibrium compaction is computed by equilibrium_compaction_si():

(6)#\[s_{eq}(t) = S_s(t)\, \Delta h(t)\, H(t)\]

with \(S_s\) (1/m), \(H\) (m), and drawdown \(\Delta h\) (m) formed from h_mean_si and h_ref_si using drawdown_rule and gated by drawdown_mode.

Parameters:
  • h_mean_si (Tensor)

  • Ss_field (Tensor)

  • H_field_si (Tensor)

  • tau_field (Tensor)

  • h_ref_si (Tensor)

  • s_init_si (Tensor)

  • dt (Any | None)

  • time_units (str | None)

  • method (str)

  • eps_tau (float)

  • relu_beta (float)

  • drawdown_mode (str)

  • drawdown_rule (str)

  • stop_grad_ref (bool)

  • drawdown_zero_at_origin (bool)

  • drawdown_clip_max (float | None)

  • verbose (int)

Return type:

Tensor

geoprior.models.subsidence.maths.compute_consolidation_step_residual(*, s_state_si, h_mean_si, Ss_field, H_field_si, tau_field, h_ref_si, dt=None, time_units='yr', method='exact', eps_tau=1e-12, relu_beta=20.0, drawdown_mode='smooth_relu', drawdown_rule='ref_minus_mean', stop_grad_ref=True, drawdown_zero_at_origin=False, drawdown_clip_max=None, verbose=0)[source]#

Compute a one-step consolidation residual in SI space.

This function forms a per-step residual that penalizes violations of a first-order consolidation relaxation model over a sequence of states. It is intended for physics diagnostics and for PDE-style training objectives where the settlement state is predicted (or derived) and should satisfy a stable time-stepping rule.

Parameters:
  • s_state_si (Tensor)

  • h_mean_si (Tensor)

  • Ss_field (Tensor)

  • H_field_si (Tensor)

  • tau_field (Tensor)

  • h_ref_si (Tensor)

  • dt (Any | None)

  • time_units (str | None)

  • method (str)

  • eps_tau (float)

  • relu_beta (float)

  • drawdown_mode (str)

  • drawdown_rule (str)

  • stop_grad_ref (bool)

  • drawdown_zero_at_origin (bool)

  • drawdown_clip_max (float | None)

  • verbose (int)

Return type:

Tensor

geoprior.models.subsidence.maths.tau_phys_from_fields(model, K_field, Ss_field, H_field, *, eps=1e-15, verbose=0, return_log=False)[source]#

Compute the physics closure consolidation timescale tau_phys and the effective drainage thickness Hd.

This function implements the model’s consolidation timescale closure \(tau_{phys}\) in a numerically stable way. The core design is to compute \(log(tau_{phys})\) first, and only apply exp at the end (unless return_log=True). This prevents unstable gradients that can arise from naive algebraic forms that contain high powers of \(1/K\).

Parameters:
  • K_field (Tensor)

  • Ss_field (Tensor)

  • H_field (Tensor)

  • eps (float)

  • verbose (int)

  • return_log (bool)

Return type:

tuple[Tensor, Tensor]

geoprior.models.subsidence.maths.compute_consistency_prior(model, K_field, Ss_field, tau_field, H_field, *, verbose=0)[source]#

Compute the consolidation timescale consistency prior.

This prior constrains the learned consolidation timescale tau to remain physically consistent with the permeability-storage- thickness closure implied by the poroelastic consolidation model. It returns the log-space mismatch:

(7)#\[R_{\mathrm{prior}} = \log(\tau_{\mathrm{learned}}) - \log(\tau_{\mathrm{phys}})\]

where \(\tau_{\mathrm{phys}}\) is computed from the predicted fields \(K\), \(S_s\), and \(H\) through tau_phys_from_fields().

Log-space is used for two reasons:

  1. Positivity: \(\tau > 0\) is enforced implicitly.

  2. Scale: timescales may span orders of magnitude; comparing logs yields a relative-type error signal.

Parameters:
  • K_field (Tensor)

  • Ss_field (Tensor)

  • tau_field (Tensor)

  • H_field (Tensor)

  • verbose (int)

Return type:

Tensor

geoprior.models.subsidence.maths.compute_smoothness_prior(dK_dx, dK_dy, dSs_dx, dSs_dy, *, K_field=None, Ss_field=None, already_log=False, verbose=0)[source]#

Compute a smoothness prior on spatial gradients of physics fields.

This function builds a spatial regularizer that penalizes rapid variation of the permeability-like field K and the storage field Ss over the spatial coordinates. In the GeoPrior PINN, this prior stabilizes the inverse problem by discouraging unrealistic high-frequency spatial structure in learned fields.

The preferred penalty is applied in log-space:

(8)#\[R_{\mathrm{smooth}} = \left\|\nabla \log K\right\|^2 + \left\|\nabla \log S_s\right\|^2\]

where, in 2D:

(9)#\[\left\|\nabla \log K\right\|^2 = \left(\frac{\partial \log K}{\partial x}\right)^2 + \left(\frac{\partial \log K}{\partial y}\right)^2\]

and similarly for \(S_s\). Penalizing gradients of logs is often preferable to raw gradients because it regularizes relative changes (order-of-magnitude variations) rather than absolute changes.

Parameters:
  • dK_dx (Tensor)

  • dK_dy (Tensor)

  • dSs_dx (Tensor)

  • dSs_dy (Tensor)

  • K_field (Any | None)

  • Ss_field (Any | None)

  • already_log (bool)

  • verbose (int)

Return type:

Tensor

geoprior.models.subsidence.maths.exp_from_bounds(raw_log, log_min, log_max, *, mode='soft', beta=20.0, guard=5.0, eps=0.0, dtype=None, name='')[source]#
geoprior.models.subsidence.maths.get_log_bounds(model, *, as_tensor=True, dtype=tf.float32, verbose=0)[source]#

Get validated log-space bounds for K and Ss.

This helper reads bounds from model.scaling_kwargs['bounds'] and returns a 4-tuple:

(logK_min, logK_max, logSs_min, logSs_max).

It supports two equivalent representations:

  • Direct log-bounds: logK_min/logK_max and logSs_min/logSs_max.

  • Linear bounds converted to logs: K_min/K_max and Ss_min/Ss_max.

If bounds are missing, the function returns (None, None, None, None).

Parameters:
  • model (Any) – Model-like object with an optional scaling_kwargs dict. Bounds are read from model.scaling_kwargs['bounds'].

  • as_tensor (bool, default True) – If True, return Tensor scalars created with tf_constant. If False, return Python floats.

  • dtype (tf.DType, default tf_float32) – Tensor dtype used when as_tensor=True.

  • verbose (int, default 0) – Verbosity level for optional debug printing.

Returns:

logK_min, logK_max, logSs_min, logSs_max – Log-space bounds as Tensor scalars (if as_tensor=True), otherwise Python floats.

If bounds are not configured, returns: (None, None, None, None).

Return type:

tuple

Raises:

ValueError

If bounds exist but are invalid, including:

  • non-finite values (NaN or inf)

  • non-positive linear bounds (<= 0)

  • unordered bounds (max <= min)

Notes

This function never emits NaN log bounds. If the configuration contains invalid entries, it fails fast with ValueError.

If log-bounds are present, they are used directly. Otherwise, the function looks for linear bounds and converts them via:

(10)#\[\log K_{\min} = \log(K_{\min}), \quad \log K_{\max} = \log(K_{\max}),\]

and similarly for \(S_s\).

If Ss_min/Ss_max appear to be compressibility-like values (e.g., \(m_v\)), the function may optionally convert them to \(S_s\) using \(S_s = m_v \gamma_w\) when a finite model.gamma_w is available. This heuristic is best-effort and never raises by itself.

Examples

Use Tensor bounds for downstream math:

>>> logK_min, logK_max, logSs_min, logSs_max = get_log_bounds(
...     model, as_tensor=True
... )

Return Python floats for inspection:

>>> bounds = get_log_bounds(model, as_tensor=False)
>>> print(bounds)

See also

get_log_tau_bounds

Companion helper for tau bounds in log space.

compute_bounds_residual

Uses these bounds to compute normalized violations.

geoprior.models.subsidence.maths.get_log_tau_bounds(model, *, as_tensor=True, dtype=tf.float32, verbose=0)[source]#

Get validated log-space bounds for the consolidation timescale.

This helper returns a 2-tuple:

(log_tau_min, log_tau_max),

where \(\tau\) is the consolidation timescale expressed in SI seconds, and the returned bounds are in log-seconds.

The function reads bounds from model.scaling_kwargs['bounds'] with the following precedence:

  1. Explicit log bounds: log_tau_min and log_tau_max (already log-seconds).

  2. Linear bounds in seconds: tau_min and tau_max.

  3. Linear bounds in dataset time units: tau_min_units and tau_max_units multiplied by the seconds-per-time-unit factor inferred from time_units.

  4. Robust defaults if nothing is provided.

Parameters:
  • model (Any) – Model-like object with an optional scaling_kwargs dict. Tau bounds are read from model.scaling_kwargs['bounds'].

  • as_tensor (bool, default True) – If True, return Tensor scalars created with tf_constant. If False, return Python floats.

  • dtype (tf.DType, default tf_float32) – Tensor dtype used when as_tensor=True.

  • verbose (int, default 0) – Verbosity level for optional debug printing.

Returns:

log_tau_min, log_tau_max – Log-space bounds (log-seconds). Returned as Tensor scalars when as_tensor=True, otherwise as Python floats.

Return type:

tuple

Raises:

ValueError

If user-provided bounds exist but are invalid, including:

  • non-finite values (NaN or inf)

  • non-positive linear bounds (<= 0)

  • unordered bounds (max <= min) for explicit log bounds

Notes

The consolidation timescale \(\tau\) controls the relaxation rate in a first-order consolidation closure, e.g.:

(11)#\[\partial_t s = \frac{s_{eq}(h) - s}{\tau},\]

where \(s\) is settlement and \(s_{eq}\) is the equilibrium settlement implied by head (or drawdown).

If no tau bounds are provided, robust defaults are used:

  • tau_min = 7 days

  • tau_max = 300 years

Both are converted to seconds and then logged. A warning may be emitted to make the defaulting explicit.

If linear bounds are provided with tau_max < tau_min, the function may swap them to maintain a valid interval.

Examples

Use Tensor bounds for log-space clipping:

>>> log_tau_min, log_tau_max = get_log_tau_bounds(model)

Return floats for reporting:

>>> log_tau_min, log_tau_max = get_log_tau_bounds(
...     model, as_tensor=False
... )

See also

get_log_bounds

Bounds helper for log(K) and log(Ss).

compute_bounds_residual

Computes normalized bound violations using these limits.

geoprior.models.subsidence.maths.bounded_exp(raw, log_min, log_max, *, eps=1e-15, return_log=False, verbose=0)[source]#

Exponentiate a raw parameter inside hard log-bounds.

This helper maps an unconstrained tensor raw to a positive field by interpolating in log space between log_min and log_max. The mapping is smooth and bounded:

(12)#\[z = \sigma(\mathrm{raw}), \quad \log v = \log v_{min} + z(\log v_{max} - \log v_{min}), \quad v = \exp(\log v) + \varepsilon,\]

where \(\sigma\) is the logistic sigmoid and \(\varepsilon\) is a small positive floor.

This is used when bounds_mode="hard" to ensure learned fields such as \(K\), \(S_s\), or \(\tau\) never leave their configured ranges.

Parameters:
  • raw (Tensor) – Unconstrained logit-like tensor (any shape). Non-finite entries are sanitized to zeros to avoid NaN propagation.

  • log_min (Tensor) – Lower bound in log space. Must be finite for strict correctness, but non-finite values are sanitized to a safe constant to prevent NaNs.

  • log_max (Tensor) – Upper bound in log space. Must be finite for strict correctness, but non-finite values are sanitized to a safe constant to prevent NaNs.

  • eps (float, default _EPSILON) – Positive floor added after exponentiation to guarantee strictly positive output.

  • return_log (bool, default False) – If True, return (out, logv) where logv is the bounded log value actually exponentiated. If False, return out only.

  • verbose (int, default 0) – Verbosity level for optional debug printing.

Returns:

  • out (Tensor) – Positive bounded field tensor with the same shape as raw.

  • logv (Tensor, optional) – Bounded log value used to compute out. Returned only when return_log=True.

Notes

The sigmoid interpolation produces values strictly inside the interval (up to numerical precision). This avoids the gradient discontinuity of direct clipping while still enforcing bounds.

To prevent NaNs and Infs from contaminating training, the function sanitizes:

  • non-finite values in raw to zeros,

  • non-finite values in bounds to safe constants,

  • swapped bounds by repairing the interval ordering.

This behavior is defensive and prioritizes numerical stability.

Examples

Bound a raw logit field to the K interval:

>>> K, logK = bounded_exp(
...     rawK, logK_min, logK_max, return_log=True
... )

Bound a tau field (already in log seconds bounds):

>>> tau = bounded_exp(raw_tau, log_tau_min, log_tau_max)

See also

guarded_exp_from_bounds

Soft-bounds path that keeps raw logs for penalties while guarding exponentiation overflow.

compose_physics_fields

Uses bounded_exp to build K, Ss, and tau fields.

geoprior.models.subsidence.maths.finite_floor(x, eps)[source]#

Replace NaN/Inf by eps and floor to eps.

Useful when you want “never NaN” behaviour, not strict errors.

Parameters:
Return type:

Tensor

geoprior.models.subsidence.maths.compose_physics_fields(model, *, coords_flat, H_si, K_base, Ss_base, tau_base, training=False, eps_KSs=1e-15, eps_tau=1e-06, verbose=0)[source]#

Compose physically meaningful fields \(K\), \(S_s\), and \(\tau\) from network “base” logits and coordinate corrections.

This routine is the central field mapping step for GeoPrior-style PINN models. The model predicts coarse (time-dependent) latent logits K_base, Ss_base, and tau_base from the physics head, then adds smooth spatial corrections from coordinate MLPs:

  • model.K_coord_mlp for \(\log K\)

  • model.Ss_coord_mlp for \(\log S_s\)

  • model.tau_coord_mlp for \(\Delta \log \tau\)

The corrected parameters are then mapped to SI-consistent, positive fields (in float32-safe ways) and combined with a physics closure timescale \(\tau_\mathrm{phys}\) computed from the fields.

Let \((t, x, y)\) denote the coordinate tensor passed to the decoder. Spatial corrections are evaluated on coordinates with time zeroed:

(13)#\[\tilde{\mathbf{c}} = (0, x, y).\]

Define the raw log-parameters (logits) as:

(14)#\[\begin{split}\ell_K &= \ell_K^\mathrm{base}(t,x,y) + \Delta \ell_K(\tilde{\mathbf{c}}), \\ \ell_{S_s} &= \ell_{S_s}^\mathrm{base}(t,x,y) + \Delta \ell_{S_s}(\tilde{\mathbf{c}}).\end{split}\]

The resulting fields are positive exponentials:

(15)#\[K = \exp(\ell_K), \qquad S_s = \exp(\ell_{S_s}),\]

subject to (log-)bounds. In bounds_mode="hard" the values are projected into the valid interval by clipping in log space, while in bounds_mode="soft" the function returns the unbounded logs for penalties but uses a guarded exponential to avoid float32 overflow.

For the consolidation timescale, we first compute a closure (prior) timescale from the fields:

(16)#\[\log \tau_\mathrm{phys} = f_\tau(K, S_s, H; \text{model options}),\]

where \(H\) is the drained thickness in meters (H_si) and tau_phys_from_fields implements the chosen closure and drainage convention. The network adds a learnable residual in log space:

(17)#\[\Delta \log \tau = \ell_\tau^\mathrm{base}(t,x,y) + \Delta \ell_\tau(\tilde{\mathbf{c}}),\]

and the total learned timescale is:

(18)#\[\log \tau = \log \tau_\mathrm{phys} + \Delta \log \tau, \qquad \tau = \exp(\log \tau) + \varepsilon_\tau.\]

The term \(\varepsilon_\tau\) (eps_tau) is a small positive floor to avoid exact zeros and improve numerical stability.

Parameters:
  • model (Any) –

    Model-like object providing:

    • coordinate MLPs: K_coord_mlp, Ss_coord_mlp, tau_coord_mlp

    • bounds configuration: bounds_mode and bounds accessors used by get_log_bounds and get_log_tau_bounds

    • closure configuration used by tau_phys_from_fields

  • coords_flat (Tensor) – Coordinate tensor used by the decoder. Expected shape is (B, H, 3) with last dimension ordered as (t, x, y). The function constructs (0, x, y) for the coordinate MLPs to keep corrections time-invariant by default.

  • H_si (Tensor) – Drained thickness \(H\) in SI units (meters). Shape must be broadcastable to (B, H, 1).

  • K_base (Tensor) – Base logits for \(\log K\). Shape is typically (B, H, 1).

  • Ss_base (Tensor) – Base logits for \(\log S_s\). Shape is typically (B, H, 1).

  • tau_base (Tensor) – Base logits for \(\Delta \log \tau\). Shape is typically (B, H, 1).

  • training (bool, default False) – Forward mode for coordinate MLPs.

  • eps_KSs (float, default _EPSILON) – Small positive constant used when mapping log-parameters to positive values (e.g., inside bounded / guarded exponentials).

  • eps_tau (float, default 1e-6) – Additive floor for \(\tau\) in seconds to avoid exact zeros.

  • verbose (int, default 0) – Verbosity level used by internal debug printing utilities.

Returns:

  • K_field (Tensor) – Effective hydraulic conductivity field \(K\) in SI units. Shape (B, H, 1). Units are typically meters per second.

  • Ss_field (Tensor) – Effective specific storage field \(S_s\) in SI units. Shape (B, H, 1). Units are typically inverse meters.

  • tau_field (Tensor) – Learned consolidation timescale \(\tau\) in seconds. Shape (B, H, 1).

  • tau_phys (Tensor) – Closure-based timescale \(\tau_\mathrm{phys}\) in seconds. Shape (B, H, 1) (broadcasted as needed).

  • Hd_eff (Tensor) – Effective drainage thickness \(H_d\) in meters used by the closure, accounting for drainage mode and hd_factor style options. Shape broadcastable to (B, H, 1).

  • delta_log_tau (Tensor) – The learnable log-residual \(\Delta \log \tau\) added to \(\log \tau_\mathrm{phys}\). Shape (B, H, 1).

  • logK (Tensor) – Log-parameter \(\log K\) used for priors, bounds penalties, and diagnostics. Shape (B, H, 1).

  • logSs (Tensor) – Log-parameter \(\log S_s\) used for priors, bounds penalties, and diagnostics. Shape (B, H, 1).

  • log_tau (Tensor) – Log of total timescale \(\log \tau\) (pre-guard in soft mode). Returned for bounds penalties and diagnostics. Shape (B, H, 1).

  • log_tau_phys (Tensor) – Log of closure timescale \(\log \tau_\mathrm{phys}\) returned for priors and diagnostics. Shape (B, H, 1).

Notes

Why coordinate corrections use ``(0, x, y)``. The coordinate MLPs are intended to represent slowly varying spatial heterogeneity (e.g., lithology-driven variability). Zeroing time reduces the risk that the model encodes time-varying physics fields that can destabilize PDE derivatives across horizons.

Hard vs soft bounds. When bounds_mode="hard", log-parameters are projected into the valid interval, yielding fields that always satisfy bounds.

When bounds_mode="soft", log-parameters are returned unmodified for differentiable penalties, but exponentiation is guarded to prevent float32 overflow. This preserves gradients for penalties without risking NaN / Inf in the forward pass.

Numerical stability. The function deliberately avoids reapplying log(exp(.)) patterns. In particular, it composes \(\log \tau\) additively:

(19)#\[\log \tau = \log \tau_\mathrm{phys} + \Delta \log \tau,\]

which is both exact and numerically stable.

Examples

Compute fields inside a physics forward pass:

>>> K_field, Ss_field, tau_field, tau_phys, Hd_eff, dlogtau, logK, \
... logSs, log_tau, log_tau_phys = compose_physics_fields(
...     model,
...     coords_flat=coords,
...     H_si=H_si,
...     K_base=K_logits,
...     Ss_base=Ss_logits,
...     tau_base=dlogtau_logits,
...     training=True,
... )

Use returned logs for priors and bounds penalties:

>>> prior_res = dlogtau
>>> bounds_penalty_inputs = (logK, logSs, log_tau)

See also

tau_phys_from_fields

Computes the closure timescale \(\tau_\mathrm{phys}\).

get_log_bounds, get_log_tau_bounds, bounded_exp, guarded_exp_from_bounds

compute_bounds_residual

Uses the returned logs and thickness for bounds penalties.

geoprior.models.subsidence.maths.compute_bounds_residual(model, *, H_field, logK=None, logSs=None, log_tau=None, K_field=None, Ss_field=None, tau_field=None, eps=1e-15, verbose=0)[source]#

Compute differentiable bound-violation residuals for the learned physics fields.

This function converts configured parameter bounds into residual maps that can be squared and averaged to form a soft penalty term (e.g., \(L_\mathrm{bounds} = \mathrm{mean}(R^2)\)).

The bounds policy is driven by model.scaling_kwargs['bounds'] and supports:

  • Linear-space bounds for drained thickness \(H\) (meters).

  • Log-space bounds for \(K\), \(S_s\), and \(\tau\).

The returned residuals are normalized by the corresponding bound ranges, so they are roughly comparable across parameters.

Parameters:
  • model (Any)

  • H_field (Tensor)

  • logK (Any | None)

  • logSs (Any | None)

  • log_tau (Any | None)

  • K_field (Any | None)

  • Ss_field (Any | None)

  • tau_field (Any | None)

  • eps (float)

  • verbose (int)

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

geoprior.models.subsidence.maths.guard_scale_with_residual(residual, scale, *, floor, eps=1e-15)[source]#

Guard a residual scale using the observed residual magnitude.

This helper prevents residual normalization from exploding when a nominal scale is too small compared with the actual residual values observed on the current batch.

Parameters:
  • residual (Tensor)

  • scale (Tensor)

  • floor (float)

  • eps (float)

Return type:

Tensor

geoprior.models.subsidence.maths.scale_residual(residual, scale, *, floor=1e-15)[source]#

Scale a residual by a (guarded) normalization factor.

This helper divides a residual tensor by a positive scale, with strict safeguards against non-finite or tiny scales. The scale is treated as a constant with respect to backpropagation (stop-gradient).

Parameters:
  • residual (Tensor)

  • scale (Tensor)

  • floor (float)

Return type:

Tensor

geoprior.models.subsidence.maths.compute_scales(model, *, t, s_mean, h_mean, K_field, Ss_field, tau_field=None, H_field=None, h_ref_si=None, Q=None, dt=None, time_units=None, dh_dt=None, div_K_grad_h=None, verbose=0)[source]#

Compute robust normalization scales for physics residuals.

This function estimates per-batch (or per-sample) scale factors used to non-dimensionalize physics residuals before squaring and averaging. The goal is to make losses comparable across sites, time spans, and coordinate encodings, and to prevent a single residual from dominating due to unit magnitude alone.

The returned scales are typically used as:

(20)#\[R_{cons}^{*} = \frac{R_{cons}}{s_{cons}}, \qquad R_{gw}^{*} = \frac{R_{gw}}{s_{gw}},\]

where \(s_{cons}\) and \(s_{gw}\) are produced by this function (with floors applied for numerical safety).

The routine is intentionally defensive. It sanitizes shapes to (B, H, 1), guards non-finite values, enforces positive dt, and applies safe floors before any division or reduction.

Parameters:
  • model (Any) –

    Model-like object holding configuration in model.scaling_kwargs and optionally model.time_units and model.h_ref. This function reads:

    • consolidation display mode from resolve_cons_units(sk)

    • groundwater display mode from resolve_gw_units(sk)

    • coordinate normalization flags via sk['coords_normalized']

    • coordinate ranges via coord_ranges(sk)

    • auto floors via resolve_auto_scale_floor(kind, sk)

  • t (Tensor) – Time coordinate tensor. Expected shape is (B, H, 1) or (B, H). Units follow the dataset time encoding. If coords_normalized=True, t is assumed normalized and is de-normalized using coord_ranges(sk)['t'] when dt is inferred internally.

  • s_mean (Tensor) – Mean settlement state used for consolidation scaling. Expected shape is (B, H, 1) or (B, H).

  • h_mean (Tensor) – Mean head state used for scaling. Expected shape is (B, H, 1) or (B, H). Units should match the model internal convention (typically SI meters).

  • K_field (Tensor) – Effective conductivity field. Present for signature compatibility and potential future scale heuristics. Current logic does not require this argument directly.

  • Ss_field (Tensor) – Effective specific storage field \(S_s\). Used by both consolidation and groundwater scale heuristics. Expected shape is broadcastable to (B, H, 1).

  • tau_field (Tensor, optional) – Consolidation timescale \(tau\) in seconds. Provide this together with H_field to enable relaxation-aware consolidation scaling.

  • H_field (Tensor, optional) – Drained thickness \(H\) in meters. Used with tau_field for relaxation-aware consolidation scaling.

  • h_ref_si (Tensor, optional) – Reference head \(h_{ref}\) in meters. If not provided, the function falls back to model.h_ref (or 0.0). The value is broadcast to (B, H, 1) and sanitized.

  • Q (Tensor, optional) – Source term used in the groundwater residual scaling. Expected shape is broadcastable to (B, H, 1).

  • dt (Tensor, optional) – Time step tensor in the dataset time units. If provided, it is used directly (after shape normalization). If None, dt is inferred from t. The inferred dt is de-normalized when coords_normalized=True.

  • time_units (str, optional) – Name of the dataset time unit (e.g., “year”, “day”, “second”). If None, the function resolves it from sk['time_units'] or model.time_units. It is used to convert dt to seconds.

  • dh_dt (Tensor, optional) – Precomputed \(dh/dt\) in SI units (m/s). If provided, groundwater scaling can use it directly rather than reconstructing a representative magnitude.

  • div_K_grad_h (Tensor, optional) – Precomputed divergence term for groundwater flow, \(\nabla \cdot (K \nabla h)\), in SI units. If provided, it is used as a representative magnitude for groundwater scaling.

  • verbose (int, default 0) – Verbosity level. If > 0, basic statistics of computed scales may be printed.

Returns:

scales – Dictionary with keys:

  • 'cons_scale' : Tensor Scale for consolidation residuals.

  • 'gw_scale' : Tensor Scale for groundwater-flow residuals.

Each value is shaped as (B, 1, 1) or broadcastable to (B, H, 1), depending on internal heuristics.

Return type:

dict[str, Tensor]

Notes

Why scaling is needed. Consolidation and groundwater residuals can differ by many orders of magnitude depending on:

  • the dataset time unit (years vs seconds),

  • coordinate normalization spans (t, x, y),

  • site geometry and hydro-mechanical priors,

  • whether residuals are reported in SI or display units.

A stable scaling strategy prevents trivial unit choices from changing optimization dynamics.

dt construction and safety. If dt is not provided, dt is inferred as consecutive differences along horizon:

  • if \(H > 1\), \(dt_h = t_{h} - t_{h-1}\)

  • else, dt defaults to 1.0 (in dataset time units)

When coords_normalized=True, dt is multiplied by the raw time span t_range from coord_ranges(sk) to recover dt in dataset time units. dt is then converted to seconds via dt_to_seconds(dt, time_units=...).

All dt paths apply:

  • absolute value

  • finite sanitization

  • a positive floor

  • a final lower bound using seconds_per_time_unit(time_units)

This guards against degenerate dt values that would explode scales.

Relaxation-aware consolidation scaling. If both tau_field and H_field are provided, consolidation scales may incorporate a relaxation time scale to better match the form of the consolidation closure used by the model. If they are not provided, a simpler heuristic is used.

Groundwater scaling inputs. Groundwater scales are computed from representative magnitudes of the groundwater PDE components, optionally using dh_dt and div_K_grad_h when provided. The scaling also accounts for display unit policies returned by resolve_gw_units(sk).

This function is not traced. This wrapper is not decorated with tf.function because it accepts a Python model object. Callers may wrap the function at a higher level if a stable tracing boundary is desired.

Examples

Compute scales inside the physics path:

>>> scales = compute_scales(
...     model,
...     t=t,
...     s_mean=s_inc_pred,
...     h_mean=h_si,
...     K_field=K_field,
...     Ss_field=Ss_field,
...     tau_field=tau_field,
...     H_field=H_si,
...     h_ref_si=h_ref_11,
...     Q=Q_si,
...     dt=dt_units,
...     time_units=model.time_units,
...     dh_dt=dh_dt,
...     div_K_grad_h=dKdhx + dKdhy,
... )

Use the returned scales to normalize residuals:

>>> cons_scaled = R_cons / scales["cons_scale"]
>>> gw_scaled = R_gw / scales["gw_scale"]

See also

scale_residual

Applies a scale and floor to a residual tensor.

resolve_auto_scale_floor

Resolves “auto” floors for scale denominators.

ensure_si_derivative_frame

Converts raw autodiff derivatives to SI-consistent forms.

geoprior.models.subsidence.maths.resolve_auto_scale_floor(key, scaling_kwargs, default_val='auto')[source]#

Robustly determine a numerical stability floor for physics scales.

If the user provides a float in scaling_kwargs, it is respected. If ‘auto’, we derive a safe floor based on float32 stability limits converted to the active unit system (SI, time_units, or steps).

Baselines (SI):
  • cons (velocity): 1e-7 m/s (~3 m/yr) High floor because velocity residuals are often noise-dominated.

  • gw (rate): 1e-9 1/s (~0.03 /yr) Lower floor to capture subtler groundwater dynamics.

Parameters:
Return type:

float

geoprior.models.subsidence.maths.resolve_gw_units(sk)[source]#
geoprior.models.subsidence.maths.resolve_cons_units(sk)[source]#

Normalize consolidation residual units.

Parameters:

sk (dict[str, Any] | None)

Return type:

str

geoprior.models.subsidence.maths.settlement_state_for_pde(s_pred_si, t, *, scaling_kwargs=None, inputs=None, time_units=None, baseline_keys=('s0_si', 'subs0_si', 's_ref_si', 'subs_ref_si'), dt=None, return_incremental=True, verbose=0)[source]#

Map predicted settlement output into a PDE-ready settlement state.

This helper converts a model settlement output s_pred_si into a consistent settlement time series in SI meters that can be used as the state variable in the consolidation residual and related physics terms.

The model can represent settlement in different output modes controlled by scaling_kwargs['subsidence_kind']:

  • "cumulative" : s_pred_si already represents cumulative settlement \(s(t)\) (meters).

  • "increment" : s_pred_si represents per-step increments \(\Delta s_h\) (meters per step).

  • "rate" : s_pred_si represents a settlement rate \(ds/dt\) (meters per time unit).

The function first constructs a cumulative series \(s(t)\) and then optionally returns the incremental state \(s_{inc}(t)\) used by the ODE/PDE residuals.

Parameters:
  • s_pred_si (Tensor) –

    Predicted settlement output in SI meters (or SI meters per time unit when subsidence_kind="rate"). Expected shapes:

    • (B, H, 1) (preferred)

    • (B, H) will be promoted to (B, H, 1)

  • t (Tensor) – Time coordinate used to infer \(\Delta t\) when subsidence_kind="rate" and dt is not provided. Expected shape is (B, H, 1) or (B, H).

  • scaling_kwargs (dict, optional) –

    Scaling and configuration dictionary. This function reads subsidence_kind via:

    get_sk(sk, 'subsidence_kind', default='cumulative')

    If not provided, defaults to {}.

  • inputs (dict[str, Tensor], optional) – Optional batch inputs that may contain a baseline settlement value \(s_0\) (SI meters). If provided, the function searches for the first available tensor among baseline_keys and uses it as \(s_0\).

  • time_units (str, optional) – Name of the dataset time unit (e.g., “year”, “day”). This argument is informational here and is logged for diagnostics. When subsidence_kind="rate", the interpretation of s_pred_si is “meters per time unit”.

  • baseline_keys (Sequence[str], default (``”s0_si”, ``"subs0_si",)

  • "s_ref_si" – Candidate keys to locate a baseline settlement tensor \(s_0\) in inputs. The first match found is used.

  • "subs_ref_si") – Candidate keys to locate a baseline settlement tensor \(s_0\) in inputs. The first match found is used.

  • dt (Tensor, optional) – Time step per horizon in dataset time units. Used only when subsidence_kind="rate". Expected shape is (B, H, 1) or (B, H). If None, dt is inferred from t by finite differences, with a fallback of 1.0 for the first step.

  • return_incremental (bool, default True) –

    If True, return the incremental settlement state:

    (21)\[s_{inc}(t_h) = s(t_h) - s_0,\]

    shaped like (B, H, 1). If False, return the cumulative settlement series \(s(t_h)\).

  • verbose (int, default 0) – Verbosity level. When > 0, prints basic diagnostics of the selected mode and intermediate tensors.

Returns:

s_state – Settlement state in SI meters with shape (B, H, 1).

If return_incremental=True the output is \(s_{inc}(t)\) (incremental since \(s_0\)). Otherwise the output is the cumulative series \(s(t)\).

Return type:

Tensor

Notes

Baseline handling. The baseline \(s_0\) is interpreted as the settlement value at the reference time \(t_0\) used by the physics residuals. If no baseline is provided, \(s_0\) defaults to zero with shape (B, 1, 1) and is broadcast over the horizon.

Cumulative construction. The function builds a cumulative settlement series \(s(t)\) according to subsidence_kind:

  1. subsidence_kind="cumulative"

    s_pred_si is assumed to already represent \(s(t)\):

    (22)#\[s(t_h) = s_{pred}(t_h).\]

    This includes cases where the caller already added a baseline, e.g., \(s(t) = s_0 + s_{inc}(t)\).

  2. subsidence_kind="increment"

    s_pred_si is interpreted as per-step increments:

    (23)#\[s(t_h) = s_0 + \sum_{j=0}^{h} \Delta s_j.\]
  3. subsidence_kind="rate"

    s_pred_si is interpreted as a rate in meters per time unit:

    (24)#\[\Delta s_h = \left(\frac{ds}{dt}\right)_h \Delta t_h, \qquad s(t_h) = s_0 + \sum_{j=0}^{h} \Delta s_j.\]

    If dt is not provided, \(\Delta t_h\) is inferred from the time coordinate tensor t using finite differences. The first step uses a fallback of 1.0 (for backward compatibility).

Incremental state for PDE/ODE residuals. Many physics residuals are written for an incremental settlement state \(s_{inc}(t)\) that starts at zero at \(t_0\). When return_incremental=True the function returns:

(25)#\[s_{inc}(t_h) = s(t_h) - s_0.\]

This makes it safe to concatenate an explicit initial state (e.g., s0_inc=0) when constructing a state sequence for an exact-step consolidation integrator.

Examples

Convert per-step increments to an incremental PDE state:

>>> sk = {"subsidence_kind": "increment"}
>>> s_inc = settlement_state_for_pde(
...     s_pred_si=ds_pred_m,
...     t=coords_t,
...     scaling_kwargs=sk,
...     inputs={"s0_si": s0_m},
...     return_incremental=True,
... )

Convert a rate output using provided dt:

>>> sk = {"subsidence_kind": "rate"}
>>> s_inc = settlement_state_for_pde(
...     s_pred_si=dsdt_pred_m_per_u,
...     t=coords_t,
...     dt=dt_units,
...     scaling_kwargs=sk,
...     inputs={"s0_si": s0_m},
...     return_incremental=True,
... )

Return the cumulative series instead:

>>> s_cum = settlement_state_for_pde(
...     s_pred_si=s_pred_m,
...     t=coords_t,
...     scaling_kwargs={"subsidence_kind": "cumulative"},
...     return_incremental=False,
... )

See also

compute_consolidation_step_residual

Builds the consolidation residual from settlement and head states.

cons_step_to_cons_residual

Converts a step residual into a residual consistent with the PDE time convention.

integrate_consolidation_mean

Integrates a consolidation mean settlement trajectory used as a physics-driven prediction path.

geoprior.models.subsidence.maths.to_rms(x, *, axis=None, keepdims=False, eps=None, ms_floor=None, rms_floor=None, nan_policy='propagate', dtype=None)[source]#

Compute the root-mean-square (RMS) of a tensor.

This utility computes:

(26)#\[\mathrm{RMS}(x) = \sqrt{\mathbb{E}[x^2]}\]

over the requested reduction axes. It is designed for robust diagnostics in physics-informed training loops, where tensors may contain extremely small values (needing float64) or occasional non-finite entries (handled via nan_policy).

Parameters:
  • x (Tensor) – Input tensor. Any shape is accepted. The computation is performed in dtype (default float32).

  • axis (int or Sequence[int] or None, default None) –

    Axis or axes to reduce.

    • If None, reduce over all dimensions and return a scalar.

    • If an int or sequence, reduce only those axes.

  • keepdims (bool, default False) – If True, keep reduced dimensions with length 1.

  • eps (float or None, default None) –

    Optional lower bound applied to the mean-square value before the square root is taken. If provided and > 0, the mean-square is floored as:

    (27)\[\mathrm{MS} = \max(\mathrm{MS}, \mathrm{eps})\]

    where \(\mathrm{MS} = \mathbb{E}[x^2]\).

  • ms_floor (float or None, default None) –

    Alias for an additional mean-square floor applied after eps. If provided and > 0, it is applied as:

    (28)\[\mathrm{MS} = \max(\mathrm{MS}, \mathrm{ms\_floor})\]

    This can be useful when you want a hard numerical floor but prefer to keep eps reserved for “epsilon-like” smoothing.

  • rms_floor (float or None, default None) –

    Optional lower bound applied after taking the square root. If provided and > 0:

    (29)\[\mathrm{RMS} = \max(\mathrm{RMS}, \mathrm{rms\_floor})\]

  • nan_policy ({"propagate", "raise", "omit"}, default "propagate") –

    Policy for handling non-finite values (NaN/Inf):

    • "propagate": Use the standard reduction. Non-finite values propagate through mean and the RMS becomes non-finite.

    • "raise": Assert that x is all finite before reducing, raising an error if NaN/Inf is present.

    • "omit": Ignore non-finite entries by treating them as missing. The RMS is computed from finite entries only:

      (30)#\[\mathrm{MS} = \frac{\sum x_i^2}{N_f}\]

      where \(N_f\) is the count of finite entries along the reduced axes (clipped to at least 1).

  • dtype (Any, default None) – Compute dtype. If None, uses tf_float32 for speed. Pass dtype=tf_float64 when diagnosing very small residuals or when accumulated rounding error matters.

Returns:

rms – RMS value reduced along axis. If axis=None the result is a scalar tensor; otherwise it has the reduced shape.

Return type:

Tensor

Notes

Flooring behavior. Floors are opt-in. If eps is None and ms_floor is None, no flooring is applied to the mean-square. If rms_floor is None, no flooring is applied to the RMS.

A common pattern for stable logging of near-zero residuals is to use a small mean-square floor with float64 diagnostics:

  • dtype=tf_float64 to reduce rounding error.

  • ms_floor to avoid taking sqrt(0) when a later operation applies log or divides by RMS.

Non-finite handling. nan_policy="omit" is intended for diagnostics and logging. For training-time physics losses, prefer cleaning tensors before the loss is computed, so gradients are well-defined.

Examples

Compute RMS over all entries:

>>> r = to_rms(x)

Compute per-batch RMS (reduce over horizon and channel axes):

>>> r_b = to_rms(x, axis=(1, 2))

Omit non-finite values when logging a residual map:

>>> eps_gw = to_rms(R_gw, nan_policy="omit", dtype=tf_float64)

Apply a small mean-square floor for stable downstream log:

>>> eps = to_rms(R, ms_floor=1e-30, dtype=tf_float64)

See also

scale_residual

Scales residuals by computed characteristic scales.

guard_scale_with_residual

Ensures a scale is safe when residuals are near zero.

geoprior.models.subsidence.maths.resolve_cons_drawdown_options(scaling_kwargs, *, default_mode='smooth_relu', default_rule='ref_minus_mean', default_stop_grad_ref=True, default_zero_at_origin=False, default_clip_max=None, default_relu_beta=20.0)[source]#

Resolve consolidation drawdown options from scaling_kwargs.

Supported keys (prefer the ‘cons_*’ names): - cons_drawdown_mode / drawdown_mode - cons_drawdown_rule / drawdown_rule - cons_stop_grad_ref / stop_grad_ref - cons_drawdown_zero_at_origin / drawdown_zero_at_origin - cons_drawdown_clip_max / drawdown_clip_max - cons_relu_beta / relu_beta

Returns:

drawdown_mode, drawdown_rule, stop_grad_ref, drawdown_zero_at_origin, drawdown_clip_max, relu_beta

Return type:

dict with keys

Parameters:
  • default_mode (str)

  • default_rule (str)

  • default_stop_grad_ref (bool)

  • default_zero_at_origin (bool)

  • default_clip_max (float | None)

  • default_relu_beta (float)

geoprior.models.subsidence.maths.normalize_time_units(u)[source]#

Normalize time unit strings.

Parameters:

u (str | None)

Return type:

str

geoprior.models.subsidence.maths.seconds_per_time_unit(time_units, *, dtype=tf.float32)[source]#

Seconds-per-unit.

Parameters:

time_units (str | None)

Return type:

Tensor

geoprior.models.subsidence.maths.ensure_3d(x)[source]#

Return a rank-3 tensor, preferring static rank when available.

Parameters:

x (Tensor)

Return type:

Tensor

geoprior.models.subsidence.maths.dt_to_seconds(dt, *, time_units)[source]#

dt(time_units) -> seconds.

Parameters:
  • dt (Tensor)

  • time_units (str | None)

Return type:

Tensor

geoprior.models.subsidence.maths.rate_to_per_second(dz_dt, *, time_units)[source]#

d/d(time_units) -> d/ds.

Parameters:
  • dz_dt (Tensor)

  • time_units (str | None)

Return type:

Tensor

geoprior.models.subsidence.maths.smooth_relu(x, *, beta=20.0)[source]#

Smooth approximation to relu(x) with controlled curvature.

Parameters:
  • x (Tensor)

  • beta (float)

Return type:

Tensor

geoprior.models.subsidence.maths.positive(x, *, eps=1e-15)[source]#

Softplus positivity.

Parameters:
Return type:

Tensor