Subsidence API reference<a class="headerlink" href="#subsidence-api-reference" title="Link to this heading">#

Notes

Step outline. This training step performs the following stages:

Unpack and canonicalize targets
Targets are normalized into a stable dict structure using _canonicalize_targets and reordered by self._order_by_output_keys. Only keys in self.output_names are retained to guarantee consistent ordering for both loss computation and logging.
Forward pass with physics precomputation
The step calls physics_core() inside a single outer GradientTape. The physics core performs its own inner tape to compute coordinate derivatives required by PDE residuals. The outer tape ensures gradients flow through both:
- supervised data predictions, and
- physics loss scalars produced by the physics pathway.
Supervised data loss
Targets are aligned to prediction shapes (including quantile layout when applicable) using _align_true_for_loss and then passed as lists to self.compiled_loss. This allows Keras to apply:
- per-output losses configured in compile,
- regularization losses in self.losses,
- sample weighting logic if configured.
Total objective
The physics loss contribution is taken from the physics bundle as physics_loss_scaled. If physics is disabled (or gated off) the contribution is treated as zero.
Gradients, scaling, and clipping
Gradients of the total objective are computed w.r.t. all trainable variables. The step then:
- applies optional parameter-specific gradient scaling via self._scale_param_grads (for example, to slow down m_v or kappa updates),
- filters NaN/Inf gradients using filter_nan_gradients,
- applies global norm clipping (default clip value is 1.0),
- applies gradients via self.optimizer.apply_gradients.
This sequence is intended to improve stability for stiff physics losses and mixed-scale parameters.
Auxiliary trackers
If the model is configured with add-on trackers (for example, quantile coverage/sharpness or other custom diagnostics), update_state is called on the supervised outputs.
Packed return
The step returns a single packed dictionary from pack_step_results() so both training logs and evaluation summaries remain consistent.

Physics loss semantics. The physics contribution returned by physics_core() is already assembled with internal multipliers and (optionally) warmup/ramp gating. In other words, physics_loss_scaled is the quantity that should be added to the supervised loss.

If you need raw components for debugging, enable physics debug options in scaling_kwargs (for example, debug_physics_grads=True) and use the debug hooks called inside this step.

Gradient sanity and debugging. This method provides multiple stability and debug mechanisms:

NaN/Inf gradient filtering before applying updates.
Global-norm clipping to limit catastrophic updates.
Optional per-term gradient checks via dbg_term_grads_finite when scaling_kwargs['debug_physics_grads'] is enabled.

These are particularly useful when PDE residuals are large early in training or when coordinate scaling is misconfigured.

Examples

Typical usage: compile and fit normally, relying on this custom train step:

>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(1e-3),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
... )
>>> history = model.fit(train_ds, validation_data=val_ds, epochs=5)

Inspect returned metrics keys during training:

>>> logs = model.train_step(next(iter(train_ds)))
>>> sorted(list(logs))[:5]
['data_loss', 'loss', 'physics_loss_scaled', 'total_loss', ...]

See also

geoprior.models.subsidence.step_core.physics_core: Shared physics pathway used to compute PDE residuals and physics loss scalars consistently across train and eval.
pack_step_results: Pack supervised metrics, physics terms, and manual trackers into a stable Keras logging dictionary.
filter_nan_gradients: Sanitize gradient lists by removing NaN/Inf tensors.
tf.clip_by_global_norm: TensorFlow utility for global-norm gradient clipping.

test_step(data)[source]#

Run one evaluation (validation/test) step for GeoPrior models.

This method overrides the standard tf.keras.Model.test_step to evaluate GeoPrior-style PINN models with dict inputs and multi-output targets. It computes:

supervised validation loss and metrics via compiled_loss and compiled metrics,
optional physics diagnostics and physics loss via _evaluate_physics_on_batch (no optimizer updates),
optional add-on tracker metrics (for example, quantile coverage and sharpness),
a unified packed logging dictionary returned by pack_step_results().

Unlike train_step(), this method does not apply gradients or update model parameters. It may still use a GradientTape internally for physics derivatives when physics is enabled, but no optimizer step occurs.

Parameters:

data (tuple) –

Keras batch payload as (inputs, targets).

inputs is a dict of tensors matching the GeoPrior input API (static, dynamic, future, coords, thickness, etc.).
targets is a dict (or dict-like) of supervised targets.

Targets are canonicalized and reordered to match self.output_names for stable loss computation.

Returns:

metrics – Dictionary of scalar tensors suitable for Keras validation logging. The exact keys depend on configured losses, metrics, and physics settings, and are produced by pack_step_results().

Typical keys include:

loss / total_loss: total evaluation objective.
data_loss: supervised loss only.
per-output losses/metrics from Keras compiled configuration.
physics summary terms (for example physics_loss_scaled, epsilons) if physics is enabled.
custom tracker metrics if add-on trackers are enabled.

Return type:

Notes

Step outline. This evaluation step follows a stable, dict-safe flow:

Unpack and canonicalize targets
Targets are normalized into a stable dict structure and reordered by output key contract.
Forward pass (supervised only)
The method calls call() via self(inputs, training=False) to obtain supervised predictions only. Aux tensors are not returned here by design.
Supervised loss and metrics
Targets are aligned to prediction shapes using _align_true_for_loss and passed to compiled_loss as ordered lists to ensure consistent behavior across Keras versions and dict wrappers.
Add-on trackers (optional)
If configured, add-on trackers are updated with targets and predictions. These trackers are purely diagnostic and do not affect loss values unless explicitly integrated elsewhere.
Physics diagnostics (optional)
If physics is enabled, the method calls _evaluate_physics_on_batch(inputs, return_maps=False) to compute physics residual summaries and a scaled physics loss.

The total evaluation objective is then:

(2)#\[L_{total} = L_{data} + L_{phys}\]

where \(L_{phys}\) is the physics loss scalar returned by the physics evaluator.

The physics evaluator may use internal autodiff to compute PDE derivatives for residual diagnostics, but gradients are not used to update parameters in test_step.
Packed return
The method returns a single packed dictionary from pack_step_results() to keep training and validation logs consistent.

When to use physics in validation. Enabling physics during validation is useful to monitor:

PDE residual RMS values (epsilon metrics),
consistency priors (for example, time-scale prior),
bounds penalties and stability signals.

If validation speed is a concern, physics can be disabled with the model physics switch (for example, _physics_off() returning True), in which case only supervised losses/metrics are computed.

Examples

Standard evaluation with physics enabled:

>>> logs = model.test_step(next(iter(val_ds)))
>>> float(logs["data_loss"])
1.23
>>> float(logs["physics_loss_scaled"])
0.01

Disable physics for faster validation (model-specific switch):

>>> model._physics_off = lambda: True
>>> logs = model.test_step(next(iter(val_ds)))
>>> "physics_loss_scaled" in logs
False  # depends on pack_step_results configuration

See also

train_step: Custom training step that computes physics loss and applies gradients.
_evaluate_physics_on_batch: Evaluation-only physics routine that computes residual diagnostics without applying optimizer updates.
pack_step_results: Pack supervised metrics, physics terms, and manual trackers into a stable Keras logging dictionary.

evaluate_physics(inputs, return_maps=False, max_batches=None, batch_size=None)[source]#

Evaluate physics diagnostics over a batch or a dataset.

This method computes physics-only diagnostics for GeoPrior-style PINN models. Supported input modes are:

a tf.data.Dataset whose scalar diagnostics are aggregated across batches;
a mapping of tensors or numpy-like arrays, optionally batched via batch_size;
a single pre-batched mapping that is evaluated once.

The returned values are intended for monitoring PDE consistency, prior adherence, and stability during training and validation.

Parameters:

inputs (dict or Dataset) –
Input payload used for physics evaluation.
- If a dict, it should follow the GeoPrior batch API and contain tensors, or array-like values when batch_size is provided.
- If a Dataset, each element should yield either an input dict or a tuple/list whose first element is the input dict.
return_maps (bool, default False) –
If True, include residual maps and learned field tensors.

In Dataset mode, maps are not aggregated across batches. The method returns maps from the last processed batch only to keep memory usage bounded and avoid ambiguous aggregation semantics.
max_batches (int or None, default None) –
Maximum number of dataset batches to process. If None, iterate through the entire dataset.

This option is useful for quick diagnostics on large datasets.
batch_size (int or None, default None) – If provided and inputs is a mapping of numpy-like arrays, wrap into a dataset and batch by this size before evaluation.

Returns:

out – Dictionary of physics diagnostics. In Dataset mode, scalar keys whose names start with 'loss_' or 'epsilon_' are aggregated by mean across processed batches. Example aggregated outputs include loss_cons, loss_gw, loss_prior, loss_smooth, loss_bounds, loss_mv, loss_q_reg, epsilon_cons, epsilon_gw, and epsilon_prior.

When return_maps=True, the output may also include maps from the last processed batch, such as residuals R_prior, R_cons, R_gw; learned fields K, Ss, tau; closure-prior fields tau_prior / tau_closure; and thickness fields H_field / H plus drainage thickness Hd. Map availability depends on the underlying physics computation and whether the batch contains the required inputs.

Return type:

dict of str to Tensor

Raises:

ValueError – If the underlying physics computation requires missing inputs (for example, thickness) or inputs have incompatible shapes.

Notes

Use this method to evaluate physics consistency independently of the supervised data loss. Typical use cases include monitoring residual RMS values, diagnosing unit or coordinate mismatches, validating bounds and priors, and generating physics maps for inspection.

This method does not compute supervised metrics. In Dataset mode, only scalar keys with loss_ or epsilon_ prefixes are aggregated across batches. Residual maps and learned fields are not aggregated; when return_maps=True, the method returns the maps from the last processed batch.

Examples

Evaluate physics scalars over a validation dataset:

>>> phys = model.evaluate_physics(val_ds, max_batches=10)
>>> float(phys["epsilon_prior"])
0.01

Evaluate physics and retrieve last-batch maps:

>>> phys = model.evaluate_physics(val_ds, return_maps=True, max_batches=1)
>>> phys["R_gw"].shape
TensorShape([B, H, 1])

Evaluate a single batch dictionary:

>>> phys = model.evaluate_physics(batch_dict, return_maps=False)
>>> sorted([k for k in phys if k.startswith("loss_")])[:3]
['loss_bounds', 'loss_cons', 'loss_gw']

Wrap numpy-like arrays into batches (mapping mode):

>>> phys = model.evaluate_physics(inputs_np, batch_size=256, max_batches=5)

See also

_evaluate_physics_on_batch: Per-batch physics diagnostics wrapper.
geoprior.models.subsidence.step_core.physics_core: Shared physics computation used for diagnostics and training.

current_mv()[source]#

Return the current value of the compressibility \(m_v\).

This is a thin convenience wrapper around _mv_value(), which handles both the trainable (log-parameterized) and fixed-scalar cases.

Returns:: Scalar tensor representing \(m_v\) in linear space.
Return type:: tf.Tensor

current_kappa()[source]#

Return the current value of the consistency coefficient \(\kappa\).

This is a thin convenience wrapper around _kappa_value(), which handles both the trainable (log-parameterized) and fixed-scalar cases.

Returns:: Scalar tensor representing \(\kappa\) in linear space.
Return type:: tf.Tensor

get_last_physics_fields()[source]#: Returns the most recent physics fields and H used by the model call. Shapes: (B, H, 1) each, matching the last forward pass.

split_data_predictions(data_tensor)[source]#

Split a combined supervised output tensor into subsidence and GWL components.

GeoPrior models often compute a single “data head” tensor whose last dimension concatenates multiple supervised targets:

(3)#\[y = [s, g]\]

where \(s\) is subsidence and \(g\) is groundwater level (or a GWL-like driver). This helper slices the last axis into:

subsidence prediction tensor s_pred
groundwater-level prediction tensor gwl_pred

The slicing is controlled by the model attributes self.output_subsidence_dim and self.output_gwl_dim.

Parameters:

data_tensor (Tensor) –

Combined supervised output tensor with last axis size output_subsidence_dim + output_gwl_dim.

Typical shapes include:

(B, H, D) for point predictions, where D = subs_dim + gwl_dim.
(B, H, Q, D) for quantile predictions. In this case, the slicing is still applied on the last dimension D.

Returns:

s_pred (Tensor) – Subsidence slice from data_tensor[..., :output_subsidence_dim].
gwl_pred (Tensor) – GWL slice from data_tensor[..., output_subsidence_dim:].

Return type:

tuple[Tensor, Tensor]

Notes

This method performs a pure tensor slice and does not apply any unit conversions. Unit handling is managed by scaling helpers elsewhere.
If quantiles are present, the Q axis is preserved and only the last axis is split.

Examples

Point outputs:

>>> y = tf.zeros([8, 3, 2])  # subs_dim=1, gwl_dim=1
>>> s_pred, gwl_pred = model.split_data_predictions(y)
>>> s_pred.shape, gwl_pred.shape
(TensorShape([8, 3, 1]), TensorShape([8, 3, 1]))

Quantile outputs:

>>> yq = tf.zeros([8, 3, 3, 2])  # (B,H,Q,D)
>>> s_pred, gwl_pred = model.split_data_predictions(yq)
>>> s_pred.shape, gwl_pred.shape
(TensorShape([8, 3, 3, 1]), TensorShape([8, 3, 3, 1]))

See also

split_physics_predictions: Split the physics-head tensor into (K, Ss, dlogtau, Q) logits.

split_physics_predictions(phys_means_raw_tensor)[source]#

Split the combined physics-head tensor into per-field logits.

GeoPrior models predict a compact “physics head” tensor whose last dimension concatenates the raw logits for multiple physics fields. This helper slices that tensor into:

K_logits : hydraulic conductivity logits
Ss_logits : specific storage logits
dlogtau_logits : relaxation time offset logits
Q_logits : optional forcing / source-term logits

The canonical ordering is:

(4)#\[p = [K, S_s, dlogtau, Q]\]

where each component is typically 1-dimensional, i.e. shape (B, H, 1) per component.

Parameters:

phys_means_raw_tensor (Tensor) –

Combined physics-head tensor. Expected shape is typically:

(B, H, P) where P is the total physics output dimension.
Some callers may supply tensors with additional axes, but the slicing always occurs along the last axis.

Returns:

K_logits (Tensor) – Slice corresponding to the conductivity logits. Shape is (..., output_K_dim) and usually (B, H, 1).
Ss_logits (Tensor) – Slice corresponding to the storage logits. Shape is (..., output_Ss_dim) and usually (B, H, 1).
dlogtau_logits (Tensor) – Slice corresponding to the relaxation-time offset logits. Shape is (..., output_tau_dim) and usually (B, H, 1).
Q_logits (Tensor) – Slice corresponding to the forcing/source logits. Shape is (..., output_Q_dim) and usually (B, H, 1).

If Q is disabled or missing from the input tensor, a zeros tensor with the appropriate broadcastable shape is returned.

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

Notes

Backward compatibility and “always return Q”. This helper is designed so downstream physics code never needs to branch on whether Q exists.

If self.output_Q_dim <= 0, Q is treated as disabled and a zeros tensor shaped like K_logits[..., :1] is returned.
If Q is enabled but phys_means_raw_tensor does not contain enough channels to include Q (older checkpoints), Q is returned as zeros with the correct shape.

This allows PDE residual code to accept a consistent signature regardless of whether Q is actually trained.

Shape and dimension conventions. The slice widths are controlled by model attributes:

output_K_dim
output_Ss_dim
output_tau_dim
output_Q_dim (optional)

If your model uses multi-dimensional physics heads, the returned tensors will preserve those widths accordingly.

Examples

Standard case with Q present:

>>> p = tf.zeros([8, 3, 4])  # [K,Ss,dlogtau,Q]
>>> K, Ss, dlogtau, Q = model.split_physics_predictions(p)
>>> K.shape, Ss.shape, dlogtau.shape, Q.shape
(TensorShape([8, 3, 1]), TensorShape([8, 3, 1]),
 TensorShape([8, 3, 1]), TensorShape([8, 3, 1]))

Backward-compatible case (no Q channel in stored tensor):

>>> p_old = tf.zeros([8, 3, 3])  # [K,Ss,dlogtau]
>>> K, Ss, dlogtau, Q = model.split_physics_predictions(p_old)
>>> Q.shape
TensorShape([8, 3, 1])

See also

compose_physics_fields: Map raw logits into bounded SI-consistent physics fields.
q_to_gw_source_term_si: Convert Q logits to the SI source term used in the GW PDE.

property lambda_offset_value: float#: Current raw value stored in the TF weight _lambda_offset.

property lambda_offset: float#

help(**kwargs)#

property mv_lr_mult: float#

Learning-rate multiplier for \(m_v\) (via log_mv).

This factor multiplies the gradient of the log-parameter log_mv inside _scale_param_grads(), allowing \(m_v\) to learn faster or slower than the rest of the network.

Returns:: Current value of the multiplier for log_mv.
Return type:: float

my_params = GeoPriorSubsNet( static_input_dim, dynamic_input_dim, future_input_dim, output_subsidence_dim=1, output_gwl_dim=1, embed_dim=32, hidden_units=64, lstm_units=64, attention_units=32, num_heads=4, dropout_rate=0.1, forecast_horizon=1, quantiles=None, max_window_size=10, memory_size=100, scales=None, multi_scale_agg='last', final_agg='last', activation='relu', use_residuals=True, use_batch_norm=False, pde_mode='both', identifiability_regime=None, mv=LearnableMV(initial_value=1e-07, trainable=True, name=learnable_mv), kappa=LearnableKappa(initial_value=1.0, trainable=True, name=learnable_kappa), gamma_w=FixedGammaW(value=9810.0, name=fixed_gamma_w, log_transform=True, non_negative=True), h_ref=FixedHRef(value=0.0, name=fixed_h_ref, log_transform=False, non_negative=False), use_effective_h=False, hd_factor=1.0, kappa_mode='kb', offset_mode='mul', bounds_mode='soft', residual_method='exact', time_units=None, use_vsn=True, vsn_units=None, mode=None, objective=None, attention_levels=None, architecture_config=None, scale_pde_residuals=True, scaling_kwargs=None, name='GeoPriorSubsNet', verbose=0 )#

property kappa_lr_mult: float#

Learning-rate multiplier for \(\kappa\) (via log_kappa).

This factor multiplies the gradient of the log-parameter log_kappa inside _scale_param_grads(), allowing \(\kappa\) to learn at a different pace than the other parameters.

Returns:: Current value of the multiplier for log_kappa.
Return type:: float

compile(lambda_cons=None, lambda_gw=None, lambda_prior=None, lambda_smooth=None, lambda_mv=None, lambda_bounds=None, lambda_q=None, lambda_offset=1.0, mv_lr_mult=1.0, kappa_lr_mult=1.0, scale_mv_with_offset=False, scale_q_with_offset=True, **kwargs)[source]#

Compile the model and configure data/physics loss weighting.

This override extends tf.keras.Model.compile() with explicit weights for each physics term used by GeoPrior PINN training, plus a global physics multiplier (lambda_offset) that can be scheduled during training.

The GeoPrior training objective (as used by train_step()) is:

(5)#\[L_{total} = L_{data} + \alpha(\text{offset_mode}, \lambda_{offset}) \, L_{phys}\]

where the physics objective is assembled from multiple components:

(6)#\[\begin{split}L_{phys} = &&\lambda_{cons} L_{cons}\\ && + \lambda_{gw} L_{gw}\\ && + \lambda_{prior} L_{prior}\\ && + \lambda_{smooth} L_{smooth}\\ && + \lambda_{mv} L_{mv}\\ && + \lambda_{bounds} L_{bounds}\\ && + \lambda_{q} L_{q}\\\end{split}\]

Each component corresponds to a residual (or penalty) computed in the shared physics core and summarized as mean-square values. The global multiplier \(alpha\) is determined by self.offset_mode:

offset_mode='mul' : \(\alpha = \lambda_{offset}\)
offset_mode='log10': \(\alpha = 10^{\lambda_{offset}}\)

The value of lambda_offset is stored in a non-trainable scalar weight self._lambda_offset (created via add_weight), which makes it safe to update during training from callbacks.

Parameters:

lambda_cons (float, default 1.0) –
Weight for the consolidation residual loss \(L_{cons}\).

This term penalizes the (scaled) consolidation residual \(R_{cons}\) derived from the settlement relaxation update, and is typically computed as:

(7)\[L_{cons} = E[ R_{cons}^2 ]\]
lambda_gw (float, default 1.0) –
Weight for the groundwater-flow residual loss \(L_{gw}\).

This term penalizes the (scaled) groundwater PDE residual \(R_{gw}\) of the form:

(8)\[R_{gw} = S_s \, \partial_t h - \nabla \cdot (K \nabla h) - Q\]

and is typically computed as:

(9)\[L_{gw} = E[ R_{gw}^2 ]\]
lambda_prior (float, default 1.0) –
Weight for the consistency prior loss \(L_{prior}\).

This term ties the learned relaxation time \(tau\) to a closure-based timescale \(tau_{phys}\) computed from the learned fields and thickness. In the current implementation the residual is commonly expressed in log space:

(10)\[R_{prior} = \log(\tau) - \log(\tau_{phys})\]

and the loss is:

(11)\[L_{prior} = E[ R_{prior}^2 ]\]
lambda_smooth (float, default 1.0) –
Weight for the smoothness prior loss \(L_{smooth}\).

This term penalizes spatial roughness in the learned hydraulic fields, typically via squared first derivatives:

(12)\[L_{smooth} = E[ (\partial_x K)^2 + (\partial_y K)^2 + (\partial_x S_s)^2 + (\partial_y S_s)^2 ]\]

It stabilizes training and encourages spatially coherent fields.
lambda_mv (float, default 0.0) –
Weight for the m_v consistency prior \(L_{mv}\).

This term is designed to provide a direct learning signal for \(m_v\) by aligning \(S_s\) with the expected relation with compressibility and water unit weight:

(13)\[S_s \approx m_v \, \gamma_w\]

A common residual is constructed in log space for stability:

(14)\[R_{mv} = \log(S_s) - \log(m_v \gamma_w)\]

and the loss is:

(15)\[L_{mv} = E[ \rho(R_{mv}) ]\]

where \(rho\) may be a robust penalty (for example, Huber) depending on scaling_kwargs configuration. When set to a positive value, this term can help constrain \(m_v\) in underdetermined settings.
lambda_bounds (float, default 0.0) –
Weight for the bounds penalty \(L_{bounds}\).

This term penalizes violations of configured parameter bounds (for example, thickness and log-parameter ranges) provided in scaling_kwargs['bounds']. When bounds_mode='soft', the penalty is differentiable and contributes to the objective:

(16)\[L_{bounds} = E[ R_{bounds}^2 ]\]

When bounds_mode='hard', parameters may be clipped or projected by the physics mapping, and this weight is typically forced to zero.
lambda_q (float, default 0.0) –
Weight for the forcing regularization term \(L_{q}\).

This term discourages excessive forcing magnitude by penalizing the mean-square of the SI source term \(Q\) used in the GW residual:

(17)\[L_{q} = E[ Q^2 ]\]

It is useful when a learnable forcing head is enabled and you want it to remain near zero unless required by data.
lambda_offset (float, default 1.0) –
Global physics multiplier stored in self._lambda_offset.

The effective multiplier applied to \(L_{phys}\) is:
- for offset_mode='mul' : \(alpha = \lambda_{offset}\)
- for offset_mode='log10': \(alpha = 10^{\lambda_{offset}}\)
self._lambda_offset is a non-trainable scalar weight so it can be updated safely during training, for example:

model._lambda_offset.assign(new_value)
mv_lr_mult (float, default 1.0) – Learning-rate multiplier applied to the gradient updates of the m_v log-parameter. This affects only the parameter update scaling, not the loss definition.
kappa_lr_mult (float, default 1.0) – Learning-rate multiplier applied to the gradient updates of the kappa log-parameter (the closure/unit-conversion factor used by the timescale prior). This affects only parameter update scaling, not the loss definition.
scale_mv_with_offset (bool, default False) –
If True, multiply the \(L_{mv}\) contribution by the global physics multiplier \(alpha\) in addition to lambda_mv.

This is useful when \(L_{mv}\) should follow the same warmup schedule as other physics terms. If False, \(L_{mv}\) is weighted only by lambda_mv.
scale_q_with_offset (bool, default True) –
If True, multiply the \(L_{q}\) contribution by the global physics multiplier \(alpha\) in addition to lambda_q.

This is commonly enabled so forcing regularization ramps in together with other physics terms during physics warmup.
kwargs (dict) – Additional keyword arguments forwarded to tf.keras.Model.compile(), such as optimizer, loss, metrics, run_eagerly, jit_compile, and so on.

Returns:

self – Returns the compiled model instance.

Return type:

Notes

Physics-off behavior. If the model physics is disabled (for example, by PDE mode settings or a physics switch), this method forces all physics weights to neutral values regardless of the inputs:

lambda_prior = 0.0
lambda_smooth = 0.0
lambda_mv = 0.0
lambda_q = 0.0
lambda_bounds = 0.0
self._lambda_offset = 1.0

This ensures that train_step() and test_step() remain stable and that logs do not contain misleading physics terms.

Validation of lambda_offset. For offset_mode='mul', lambda_offset must be strictly positive. For offset_mode='log10', any real value is allowed and acts as a log10-scale controller.

Scheduling lambda_offset. A recommended pattern is to keep individual lambda_* values fixed and schedule lambda_offset (warmup/ramp) using a callback. Because self._lambda_offset is a non-trainable TF weight, it is safe to update at runtime.

Examples

Compile with physics enabled and a moderate prior:

>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(1e-3),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
...     lambda_cons=1.0,
...     lambda_gw=1.0,
...     lambda_prior=2.0,
...     lambda_smooth=0.1,
...     lambda_bounds=0.01,
...     lambda_offset=0.1,
... )

Disable forcing penalty and use a stronger smoothness prior:

>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(5e-4),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
...     lambda_q=0.0,
...     lambda_smooth=1.0,
... )

Use log10 scaling for the global physics multiplier:

>>> model.offset_mode = "log10"
>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(1e-3),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
...     lambda_offset=-1.0,  # physics multiplier = 0.1
... )

See also

train_step: Uses the configured lambdas to assemble the total loss and apply gradients.
_physics_loss_multiplier: Computes the global physics multiplier from offset_mode and self._lambda_offset.
geoprior.models.subsidence.step_core.physics_core: Computes per-batch physics residuals and loss terms.

export_physics_payload(dataset, max_batches=None, save_path=None, format='npz', overwrite=False, metadata=None, random_subsample=None, float_dtype=<class 'numpy.float32'>, log_fn=None, **tqdm_kws)[source]#

Export physics diagnostics as a flat payload.

This helper collects physics diagnostics from a trained GeoPrior-style model and optionally persists them to disk.

Internally, it calls gather_physics_payload() to iterate over dataset and evaluate physics maps and scalar summaries via GeoPriorSubsNet.evaluate_physics() with return_maps=True. The per-batch tensors are flattened and concatenated into 1D arrays suitable for scatter plots, histograms, and reproducibility artifacts.

Parameters:

dataset (iterable) – Batched iterable (typically a tf.data.Dataset) yielding either inputs or (inputs, targets). Targets, if present, are ignored. Each inputs must contain the tensors required by evaluate_physics() (notably the coordinate tensor and thickness field, depending on the model configuration).
max_batches (int or None, default None) – Maximum number of batches to process. If None, consumes the entire iterable.
save_path (str or None, default None) – If provided, write the payload to this location using save_physics_payload(). If save_path is a directory, a default filename is used by the saver.
format ({'npz', 'csv', 'parquet'}, default 'npz') – Output format for persistence. 'npz' writes a compressed NumPy archive and a JSON sidecar metadata file.
overwrite (bool, default False) – If False and save_path already exists, raise an error.
metadata (dict or None, default None) – Optional user metadata to merge into the auto-generated provenance returned by default_meta_from_model(). User keys override defaults on conflict.
random_subsample (float or None, default None) – If provided, randomly subsample the flat payload after it is gathered. Must be in (0, 1] and is interpreted as the fraction of rows to keep. This is useful to reduce file size for large grids.
float_dtype (numpy dtype, default numpy.float32) – Dtype used when casting flattened arrays. Using float32 keeps files compact and is typically sufficient for diagnostics.
log_fn (callable or None, default None) – Optional logger used by the progress helper (for example, print). If None, the progress helper may be silent.
**tqdm_kws – Extra keyword arguments forwarded to the progress helper used inside gather_physics_payload().

Returns:

payload – Flat diagnostics payload with 1D arrays. The exact keys are defined by gather_physics_payload(), but typically include:

tau : effective relaxation time (seconds)
tau_prior / tau_closure : closure timescale (seconds)
K : effective hydraulic conductivity (m/s)
Ss : effective specific storage (1/m)
Hd : effective drainage thickness (m)
cons_res_vals : consolidation residual values
log10_tau and log10_tau_prior
metrics : nested dict with summary scalars

Return type:

dict[str, numpy.ndarray]

Notes

This routine does not change units. Unit consistency is a responsibility of the model physics and its scaling_kwargs.
If return_maps=True is used inside evaluate_physics(), maps are collected per batch and then flattened here. When saving, the payload is stored exactly as returned by the model.
Random subsampling is performed after concatenation, so it samples rows uniformly across all processed batches.

See also

gather_physics_payload: Core collector that builds the flat arrays.
save_physics_payload: Persist payload + metadata to disk.
default_meta_from_model: Build lightweight provenance metadata from a model.
GeoPriorSubsNet.evaluate_physics: Compute physics scalars and (optionally) maps.

Examples

>>> # ds is a batched tf.data.Dataset yielding (inputs, targets)
>>> payload = model.export_physics_payload(
...     ds, max_batches=20, random_subsample=0.25
... )
>>> # Save to disk (creates a .meta.json sidecar for npz/csv/parquet)
>>> _ = model.export_physics_payload(
...     ds,
...     max_batches=50,
...     save_path="physics_payload.npz",
...     format="npz",
...     overwrite=True,
... )

static load_physics_payload(path)[source]#

Load a previously saved physics payload.

This is a thin convenience wrapper around load_physics_payload() from the diagnostics payload module. It reads the data file and its optional JSON sidecar metadata.

Parameters:

path (str) – Path to a saved payload. Supported extensions depend on the underlying loader and typically include .npz, .csv, and .parquet. For formats that support it, a sidecar metadata file is expected at path + '.meta.json'.

Returns:

(payload, meta) –

payloaddict[str, numpy.ndarray]: Dictionary of arrays loaded from disk. Backward- and forward-compatible aliases may be added by the loader (for example, ensuring both tau_prior and tau_closure are present).
metadict: Metadata dictionary loaded from the JSON sidecar if found, otherwise an empty dict.

Return type:

tuple(dict, dict)

Notes

This method performs I/O only. It does not validate that the payload matches a particular model instance.
If you saved with format='npz', the payload is loaded using NumPy. For CSV/Parquet, the loader typically uses pandas.

See also

load_physics_payload: The underlying loader that performs format dispatch.
GeoPriorSubsNet.export_physics_payload: Export and optionally save a payload.

Examples

>>> payload, meta = GeoPriorSubsNet.load_physics_payload(
...     "physics_payload.npz"
... )
>>> list(payload)[:5]
['tau', 'tau_prior', 'K', 'Ss', 'Hd']

get_config()[source]#

Return a Keras-serializable configuration for model reconstruction.

This method extends tf.keras.Model.get_config() to ensure GeoPriorSubsNet can be saved and reloaded with tf.keras.models.load_model() (or keras.models.load_model()) while preserving the model’s physics options and scaling pipeline.

The returned dictionary contains:

the base configuration from BaseAttentive (via super().get_config()),
the supervised output layout (output_dim),
the resolved scaling configuration serialized as a Keras object,
GeoPrior-specific physics constructor arguments and flags.

The output is designed to be JSON-serializable by Keras. Objects that are not plain JSON (for example, GeoPriorScalingConfig and scalar wrappers such as LearnableMV) are included as Keras serialized objects via keras.saving.serialize_keras_object().

Returns:: config – A configuration dictionary that can be passed to from_config() to reconstruct the model.
Return type:: dict

Notes

output_dim is kept for compatibility with the BaseAttentive constructor signature. It is not a user-facing argument for the GeoPrior model; it is derived from:

(18)#\[output\_dim = output\_subsidence\_dim + output\_gwl\_dim\]
scaling_kwargs is stored as a serialized Keras object representing the validated scaling configuration. This preserves the exact conventions (units, coordinate normalization, bounds) used during training and is critical for consistent inference.
This config does not include runtime-only state such as optimizer variables or training metrics. Those are handled by standard Keras checkpointing mechanisms.

Examples

Serialize and reconstruct manually:

>>> cfg = model.get_config()
>>> model2 = model.__class__.from_config(cfg)

Save and reload through Keras:

>>> model.save("geoprior.keras")
>>> model2 = keras.models.load_model(
...     "geoprior.keras",
...     custom_objects={"GeoPriorSubsNet": GeoPriorSubsNet},
... )

See also

from_config: Reconstruct a model instance from the serialized config.
keras.saving.serialize_keras_object: Keras helper used to serialize non-JSON config objects.

classmethod from_config(config, custom_objects=None)[source]#

Rebuild a GeoPrior model instance from a serialized configuration.

This classmethod reconstructs the model from a configuration dictionary produced by get_config() and used by the Keras serialization stack.

The method performs three reconstruction steps:

Build a custom_objects registry that includes all GeoPrior wrappers and scaling configuration classes needed for safe deserialization.
Rehydrate wrapper objects stored as Keras-serialized dicts ({"class_name": ..., "config": ...}) for keys such as mv, kappa, gamma_w, and h_ref.
Rehydrate the scaling configuration stored under scaling_kwargs if present as a Keras object.

Finally, the method removes legacy/internal keys that are not part of the current constructor signature and returns cls(**config).

Parameters:

config (dict) – Serialized configuration dictionary. Typically produced by get_config() and passed by Keras during deserialization.
custom_objects (dict or None, default None) – Optional mapping used by Keras to resolve custom layers, models, and config objects. If None, an internal registry is created and merged with any user-provided entries.

Returns:

model – A reconstructed model instance equivalent to the original model at save time (architecture and configuration). Weights are loaded by Keras separately when using keras.models.load_model().

Return type:

Notes

This method is designed to be robust to older saved configs by explicitly dropping keys that were used by previous GeoPrior/PINN variants (for example, legacy groundwater coefficient keys and internal version markers).
The deserialization process relies on Keras helpers and the custom_objects registry. If you have custom subclasses or external layers referenced inside architecture_config, you must provide them in custom_objects or register them with Keras before loading.
If scaling deserialization fails, the method raises the underlying exception because the scaling configuration is required for consistent unit handling and PDE residual computation.

Examples

Reconstruct from a saved config dictionary:

>>> cfg = model.get_config()
>>> model2 = GeoPriorSubsNet.from_config(
...     cfg,
...     custom_objects={"GeoPriorSubsNet": GeoPriorSubsNet},
... )

Load a saved model with explicit custom_objects:

>>> model2 = keras.models.load_model(
...     "geoprior.keras",
...     custom_objects={
...         "GeoPriorSubsNet": GeoPriorSubsNet,
...         "GeoPriorScalingConfig": GeoPriorScalingConfig,
...     },
... )

See also

get_config: Produce the configuration dictionary used for reconstruction.
keras.saving.deserialize_keras_object: Keras helper used to rehydrate serialized config objects.

class geoprior.models.subsidence.PoroElasticSubsNet(*args, **kwargs)[source]#

Bases: GeoPriorSubsNet

Poroelastic surrogate variant of GeoPriorSubsNet.

This model is architecturally identical to GeoPriorSubsNet and follows the same dict-input API, outputs, and parameter semantics. It is provided as a physics-driven baseline for ablation and comparison runs.

Parameters:

static_input_dim (int)
dynamic_input_dim (int)
future_input_dim (int)
pde_mode (str)
use_effective_h (bool)
hd_factor (float)
kappa_mode (str)
scale_pde_residuals (bool)
scaling_kwargs (dict[str, Any] | None)
name (str)

help(**kwargs)#

my_params = PoroElasticSubsNet( static_input_dim, dynamic_input_dim, future_input_dim, pde_mode='consolidation', use_effective_h=True, hd_factor=0.6, kappa_mode='bar', scale_pde_residuals=True, scaling_kwargs=None, name='PoroElasticSubsNet' )#

__init__(static_input_dim, dynamic_input_dim, future_input_dim, pde_mode='consolidation', use_effective_h=True, hd_factor=0.6, kappa_mode='bar', scale_pde_residuals=True, scaling_kwargs=None, name='PoroElasticSubsNet', **kwargs)[source]#

Parameters:

static_input_dim (int)
dynamic_input_dim (int)
future_input_dim (int)
pde_mode (str)
use_effective_h (bool)
hd_factor (float)
kappa_mode (str)
scale_pde_residuals (bool)
scaling_kwargs (dict[str, Any] | None)
name (str)

compile(lambda_cons=1.0, lambda_gw=0.0, lambda_prior=5.0, lambda_smooth=1.0, lambda_mv=0.1, lambda_bounds=0.05, mv_lr_mult=0.5, kappa_lr_mult=0.5, **kwargs)[source]#

Compile with stronger defaults for the geomechanical prior.

Compared to GeoPriorSubsNet, this variant:

sets lambda_gw=0.0 (no groundwater-flow residual),
increases lambda_prior and lambda_bounds so that \(tau\) is tightly tied to \(tau_phys\),
gives \(m_v\) and \(kappa\) a smaller LR multiplier so they move more conservatively.

Parameters:

lambda_cons (float)
lambda_gw (float)
lambda_prior (float)
lambda_smooth (float)
lambda_mv (float)
lambda_bounds (float)
mv_lr_mult (float)
kappa_lr_mult (float)

geoprior.models.subsidence.finalize_scaling_kwargs(sk)[source]#

Add derived SI conversion constants to scaling_kwargs.

Adds (when possible): - seconds_per_time_unit: float - coord_ranges_si: dict with keys t (seconds), x/y (meters) - coord_inv_ranges_si: inverse of the above (safe floor).

Notes

This helper is designed to be called once when assembling scaling_kwargs (e.g., in your stage2 script) so the model can reuse those constants without recomputing unit conversions in the hot training loop.

Parameters:: sk (dict[str, Any])
Return type:: dict[str, Any]

geoprior.models.subsidence.debug_model_reload(mem_model, load_model, dataset, *, pred_key='subs_pred', also_check=None, top_weights=30, atol=1e-06, rtol=1e-06, log_fn=None)[source]#

Run a compact reload debug on one batch and return a dict report.

Compares predictions (max/mean abs diff) for pred_key (+ optional keys).
Compares weights by name (MISSING/SHAPE/OK).
Compares scaling_kwargs digest + time_units attribute.

Parameters:

mem_model (Any)
load_model (Any)
dataset (Iterable)
pred_key (str)
also_check (list[str] | None)
top_weights (int)
atol (float)
rtol (float)
log_fn (Callable[[str], None] | None)

Return type:

geoprior.models.subsidence.autoplot_geoprior_history(history, *, outdir, prefix='geoprior', style='default', log_fn=None)[source]#

Parameters:

history (History | dict)
outdir (str)
prefix (str)
style (str)
log_fn (Callable[[...], None] | None)

Return type:

None

geoprior.models.subsidence.plot_physics_values_in(payload, *, keys=None, dataset=None, coords=None, mode='map', title='Physics diagnostics', n_cols=2, figsize=None, savefig=None, show=True, clip_q=(0.01, 0.99), transform=None, bins=80, s=8, log_fn=None, **scatter_kwargs)[source]#: Plot physics arrays (residuals/fields) from a payload dict.

geoprior.models.subsidence.load_physics_payload(path)[source]#

Load a previously saved physics payload and its metadata.

Parameters:: path (str) – Data file path. Supports .npz, .csv, .parquet.
Returns:: (payload, meta) – Payload dict with arrays and metadata dict (if found).
Return type:: (dict, dict)

geoprior.models.subsidence.override_scaling_kwargs(sk, cfg, *, finalize=None, dyn_names=None, gwl_dyn_index=None, base_dir=None, path_key='SCALING_KWARGS_JSON_PATH', strict=True, add_path=True, log_fn=None)[source]#

Override scaling_kwargs from a JSON file or dict.

This helper applies an optional, precedence-based override to an existing scaling_kwargs mapping. The override source is read from cfg[path_key]. If the key is missing or empty, the input sk is returned (optionally finalized).

The override can be provided as:

a file path to a JSON object (mapping), or
a Python dict-like mapping embedded in cfg.

Overrides are applied via a deep-merge strategy:

for nested dict values, keys are merged recursively,
for non-dict values, the override replaces the base value.

Optionally, the merged result is passed through finalize to recompute derived or canonical fields (for example, coordinate ranges, unit flags, or other normalization metadata).

Parameters:

sk (Mapping[str, Any]) – Base scaling configuration (scaling_kwargs). This is typically computed by Stage-2 or loaded from Stage-1 output. The input is copied to a plain dict before modification.
cfg (Mapping[str, Any] or None) – Configuration mapping that may contain the override source under path_key. If None, no override is applied.
finalize (callable or None, optional) –
Function applied to the scaling dict to enforce canonical structure or to compute derived fields. If provided, it is applied before and after the override merge:
- pre-merge: normalize the base dict,
- post-merge: ensure the merged dict is consistent.
The callable must accept a dict and return a dict.
dyn_names (Sequence[str] or None, optional) – Expected dynamic feature names for safety validation. If provided and the override contains dynamic_feature_names, the two sequences are compared. A mismatch raises an error when strict=True.
gwl_dyn_index (int or None, optional) – Expected dynamic index for the groundwater-level feature. If provided and the override contains gwl_dyn_index, the values are compared. A mismatch raises an error when strict=True.
base_dir (str or None, optional) – Base directory used to resolve relative JSON paths. If None, the current working directory is used.
path_key (str, default "SCALING_KWARGS_JSON_PATH") – Name of the key in cfg that specifies the override. The value may be a dict-like mapping or a path to a JSON file.
strict (bool, default True) – Controls behavior on safety-check mismatches. When True, mismatches raise a ValueError. When False, mismatches can be logged via log_fn and the override still proceeds.
add_path (bool, default True) – If True, store the resolved override source in the output dict under scaling_kwargs_override_path. When the override is provided as a mapping (not a file), the value is set to "<dict>".
log_fn (callable or None, optional) – Optional logger function. If provided, it is called with informative messages such as successful override application and (when strict=False) mismatch warnings. Common choices are print or logger.info.

Returns:

out – Final scaling dict after optional override and optional finalization. The returned dict is independent from the input mapping object sk (a copy is always created).

Return type:

Raises:

FileNotFoundError – If cfg[path_key] is a path and the file does not exist.
ValueError – If a path is provided but the file does not contain valid JSON, or if a safety check fails while strict=True.
TypeError – If the loaded override is not a JSON object (dict-like).

Notes

Path resolution

When cfg[path_key] is a string path, it is resolved as:

Expand environment variables and ~.
If relative, join with base_dir (or CWD).

Safety checks

The checks are intentionally conservative. They prevent using an override file produced for a different dataset or feature layout. Recommended checks are:

dynamic_feature_names equality when known.
gwl_dyn_index equality when known.

You can extend validation by checking additional keys such as coord_epsg_used, coords_normalized, or unit flags.

Finalization In GeoPrior pipelines, finalize is typically a helper that enforces defaults and recomputes derived entries. Applying it both before and after the override helps reduce edge cases where the override only supplies partial information.

Figure assembly follows the plotting conventions described in Hunter [15].

Examples

Stage-2: override computed scaling with a file

In Stage-2, call this right after the auto-computed scaling is available, so the override takes precedence:

>>> sk = subsmodel_params["scaling_kwargs"]
>>> sk = override_scaling_kwargs(
...     sk,
...     cfg,
...     finalize=finalize_scaling_kwargs,
...     dyn_names=DYN_NAMES,
...     gwl_dyn_index=GWL_DYN_INDEX,
...     base_dir=os.path.dirname(__file__),
...     strict=True,
...     log_fn=print,
... )
>>> subsmodel_params["scaling_kwargs"] = sk

Stage-3: override Stage-1 scaling prior to enforcing bounds

In Stage-3, apply the override before injecting Stage-3 bounds:

>>> sk_model = dict(cfg.get("scaling_kwargs", {}) or {})
>>> sk_model = override_scaling_kwargs(
...     sk_model,
...     cfg,
...     dyn_names=sk_model.get("dynamic_feature_names"),
...     gwl_dyn_index=sk_model.get("gwl_dyn_index"),
...     base_dir=os.path.dirname(__file__),
... )
>>> sk_model["bounds"] = {
...     **(sk_model.get("bounds", {}) or {}),
...     **bounds_for_scaling,
... }

Inline dict override (no JSON file)

If the override is embedded in config, it is used directly:

>>> cfg = {
...     "SCALING_KWARGS_JSON_PATH": {
...         "coords_normalized": True,
...         "coord_ranges": {"t": 7.0, "x": 1000.0, "y": 900.0},
...     }
... }
>>> out = override_scaling_kwargs({}, cfg)

See also

finalize_scaling_kwargs: Canonicalize and complete scaling_kwargs entries.
compute_scaling_kwargs: Build a base scaling dict from data and pipeline settings.

Module index#

The index below uses fully qualified import paths throughout. It deliberately documents the submodules explicitly and avoids package-relative lookup.

The package itself is documented in the automodule block above, so it is not repeated in the autosummary tables below. Keeping the package out of the autosummary index makes stub generation more predictable and avoids self-referential package entries.

Core scientific modules#

`geoprior.models.subsidence.models`	Subsidence PINN models
`geoprior.models.subsidence.maths`	GeoPrior maths helpers (physics terms + scaling).
`geoprior.models.subsidence.scaling`	GeoPrior scaling config helpers (Keras-serializable).
`geoprior.models.subsidence.losses`	GeoPrior loss assembly and logging helpers.
`geoprior.models.subsidence.identifiability`	Identifiability scenarios for GeoPrior-style models.
`geoprior.models.subsidence.payloads`	Physics diagnostics payloads.
`geoprior.models.subsidence.utils`	GeoPrior subsidence model utilities.
`geoprior.models.subsidence.stability`	Numerical stability helpers for subsidence physics workflows.
`geoprior.models.subsidence.step_core`	Core step computations for subsidence physics evaluation.

Supporting scientific and diagnostics modules#

`geoprior.models.subsidence.batch_io`	Batch.io
`geoprior.models.subsidence.debugs`	Debug helpers for GeoPriorSubsNet.
`geoprior.models.subsidence.derivatives`	Derivative helpers for GeoPrior PINN blocks.
`geoprior.models.subsidence.doc`	Shared documentation fragments for GeoPrior PINN models.
`geoprior.models.subsidence.log_offsets_diagnostics`	Diagnostics for subsidence log-offset policies and payloads.
`geoprior.models.subsidence.plot`	Plotting helpers for subsidence training and diagnostics.

Core model classes#

The most important public exports in this package are the two model classes:

GeoPriorSubsNet
PoroElasticSubsNet

These are the main entry points for physics-guided subsidence forecasting in GeoPrior-v3.

Subsidence PINN models

class geoprior.models.subsidence.models.GeoPriorSubsNet(*args, **kwargs)[source]

Bases: BaseAttentive

Prior-regularized physics-informed network for multi-step subsidence forecasting with groundwater coupling.

GeoPriorSubsNet combines a BaseAttentive encoder-decoder with a set of physics losses that constrain the forecast to respect a simplified groundwater-flow equation and a consolidation closure. In addition, it learns spatially varying physics fields and regularizes them against geologically motivated priors.

Parameters:

static_input_dim (int)
dynamic_input_dim (int)
future_input_dim (int)
output_subsidence_dim (int)
output_gwl_dim (int)
embed_dim (int)
hidden_units (int)
lstm_units (int)
attention_units (int)
num_heads (int)
dropout_rate (float)
forecast_horizon (int)
quantiles (list[float] | None)
max_window_size (int)
memory_size (int)
scales (list[int] | None)
multi_scale_agg (str)
final_agg (str)
activation (str)
use_residuals (bool)
use_batch_norm (bool)
pde_mode (str | list[str])
identifiability_regime (str | None)
mv (LearnableMV | float)
kappa (LearnableKappa | float)
gamma_w (FixedGammaW | float)
h_ref (FixedHRef | float | str | None)
use_effective_h (bool)
hd_factor (float)
kappa_mode (str)
offset_mode (str)
bounds_mode (str)
residual_method (str)
time_units (str | None)
use_vsn (bool)
vsn_units (int | None)
mode (str | None)
objective (str | None)
attention_levels (str | list[str] | None)
architecture_config (dict | None)
scale_pde_residuals (bool)
scaling_kwargs (dict[str, Any] | None)
name (str)
verbose (int)

OUTPUT_KEYS = ('subs_pred', 'gwl_pred')

__init__(static_input_dim, dynamic_input_dim, future_input_dim, output_subsidence_dim=1, output_gwl_dim=1, embed_dim=32, hidden_units=64, lstm_units=64, attention_units=32, num_heads=4, dropout_rate=0.1, forecast_horizon=1, quantiles=None, max_window_size=10, memory_size=100, scales=None, multi_scale_agg='last', final_agg='last', activation='relu', use_residuals=True, use_batch_norm=False, pde_mode='both', identifiability_regime=None, mv=LearnableMV(initial_value=1e-07, trainable=True, name=learnable_mv), kappa=LearnableKappa(initial_value=1.0, trainable=True, name=learnable_kappa), gamma_w=FixedGammaW(value=9810.0, name=fixed_gamma_w, log_transform=True, non_negative=True), h_ref=FixedHRef(value=0.0, name=fixed_h_ref, log_transform=False, non_negative=False), use_effective_h=False, hd_factor=1.0, kappa_mode='kb', offset_mode='mul', bounds_mode='soft', residual_method='exact', time_units=None, use_vsn=True, vsn_units=None, mode=None, objective=None, attention_levels=None, architecture_config=None, scale_pde_residuals=True, scaling_kwargs=None, name='GeoPriorSubsNet', verbose=0, **kwargs)[source]

Parameters:

static_input_dim (int)
dynamic_input_dim (int)
future_input_dim (int)
output_subsidence_dim (int)
output_gwl_dim (int)
embed_dim (int)
hidden_units (int)
lstm_units (int)
attention_units (int)
num_heads (int)
dropout_rate (float)
forecast_horizon (int)
quantiles (list[float] | None)
max_window_size (int)
memory_size (int)
scales (list[int] | None)
multi_scale_agg (str)
final_agg (str)
activation (str)
use_residuals (bool)
use_batch_norm (bool)
pde_mode (str | list[str])
identifiability_regime (str | None)
mv (LearnableMV | float)
kappa (LearnableKappa | float)
gamma_w (FixedGammaW | float)
h_ref (FixedHRef | float | str | None)
use_effective_h (bool)
hd_factor (float)
kappa_mode (str)
offset_mode (str)
bounds_mode (str)
residual_method (str)
time_units (str | None)
use_vsn (bool)
vsn_units (int | None)
mode (str | None)
objective (str | None)
attention_levels (str | list[str] | None)
architecture_config (dict | None)
scale_pde_residuals (bool)
scaling_kwargs (dict[str, Any] | None)
name (str)
verbose (int)

build(input_shape)[source]

Build the model’s weights and sublayers.

Keras may call build() (e.g. via model.build() or model.summary()) before the first forward pass. For subclassed models, we must ensure all sublayers are actually built, otherwise Keras can mark the layer as built while internal state remains unbuilt.

Parameters:: input_shape (Any)
Return type:: None

property metrics: List of all metrics.

run_encoder_decoder_core(static_input, dynamic_input, future_input, coords_input, training)[source]

Run the shared encoder-decoder core for GeoPrior inputs.

This override keeps the coordinate tensor aligned with the learned sequence features that are later consumed by the physics stack.

Parameters:

static_input (Tensor)
dynamic_input (Tensor)
future_input (Tensor)
coords_input (Tensor)
training (bool)

Return type:

tuple[Tensor, Tensor]

forward_with_aux(inputs, training=False)[source]

Return predictions and auxiliary tensors for diagnostics.

This method is a thin, public wrapper around _forward_all() that exposes both:

y_pred: the supervised outputs (what call() returns),
aux: intermediate tensors useful for debugging, physics evaluation, and research diagnostics.

Unlike call(), this method is intended for inspection and tooling. It does not change Keras training behavior because it does not alter loss computation or variable updates; it simply returns additional tensors already produced by the internal forward path.

Parameters:

inputs (dict) –
Dict-input batch compatible with GeoPrior PINN models.

Typical entries include:
- static_features : Tensor, shape (B, S)
- dynamic_features : Tensor, shape (B, H, D)
- future_features : Tensor, shape (B, H, F)
- coords : Tensor, shape (B, H, 3) with last axis ordered as (t, x, y)
- H_field or soil_thickness : Tensor, thickness field broadcastable to (B, H, 1)
The exact required keys depend on the model configuration and Stage-1 export. This wrapper delegates all parsing and validation to _forward_all().
training (bool, default False) – Forward-pass training flag. When True, dropout, batch norm, and other training-time layers behave accordingly.

Returns:

y_pred (dict of str to Tensor) – Supervised predictions in the same format as call(). At minimum, keys include 'subs_pred' and 'gwl_pred'.
aux (dict of str to Tensor) – Auxiliary tensors for diagnostics. Typical keys include:
- data_final: final data head tensor used for supervised outputs (may include quantile axis).
- data_mean_raw: mean-path output before quantile modeling.
- phys_mean_raw: concatenated physics logits (K, Ss, dlogtau, optional Q).
- phys_features_raw_3d: physics feature tensor emitted by the shared encoder-decoder core.

Return type:

tuple[dict[str, Tensor], dict[str, Tensor]]

Notes

This method is recommended for:

debugging NaN/Inf propagation (by inspecting aux),
computing physics residuals outside train_step using the same forward tensors,
building evaluation utilities that need intermediate heads.

Examples

Run a forward pass and inspect physics logits:

>>> y_pred, aux = model.forward_with_aux(batch, training=False)
>>> aux["phys_mean_raw"].shape
TensorShape([B, H, 4])

See also

call: Standard Keras forward that returns supervised outputs only.
_forward_all: Internal forward routine that returns both predictions and auxiliary tensors.

call(inputs, training=False)[source]

Keras forward method returning supervised outputs only.

This method defines the standard inference and training forward behavior expected by tf.keras.Model. It returns only the supervised output dictionary that participates in Keras loss computation and metric updates.

Internally, call() delegates to _forward_all() and discards the auxiliary outputs to ensure a stable, minimal prediction contract.

Parameters:

inputs (dict) –
Dict-input batch compatible with GeoPrior PINN models.

Typical entries include:
- static_features : Tensor, shape (B, S)
- dynamic_features : Tensor, shape (B, H, D)
- future_features : Tensor, shape (B, H, F)
- coords : Tensor, shape (B, H, 3) with last axis ordered as (t, x, y)
- H_field or soil_thickness : Tensor, thickness field
All parsing, shape checks, and coordinate handling are performed by _forward_all().
training (bool, default False) – Forward-pass training flag. When True, training-time behavior (dropout, batch norm, etc.) is enabled.

Returns:

y_pred – Supervised prediction dictionary. Keys are ordered by the model output contract (for example, ('subs_pred', 'gwl_pred')). Each tensor is typically shaped:

without quantiles: (B, H, 1)
with quantiles: (B, H, Q, 1) or a model-defined quantile layout

Return type:

dict of str to Tensor

Notes

Auxiliary tensors such as physics logits and intermediate features are intentionally excluded from the return value. Use forward_with_aux() when diagnostics are required.

Examples

Standard inference call:

>>> y = model(batch, training=False)
>>> sorted(y.keys())
['gwl_pred', 'subs_pred']

See also

forward_with_aux: Forward wrapper returning both predictions and diagnostics.
_forward_all: Internal routine returning (y_pred, aux).

train_step(data)[source]

Run one custom training step for GeoPrior-style PINN training.

This method overrides the standard tf.keras.Model.train_step to train a hybrid, physics-informed model with dict inputs and multi-output supervision. The step integrates:

supervised data losses (from compile / compiled_loss),
physics losses computed by physics_core(),
optional gradient scaling for selected parameters,
robust gradient sanitization and global-norm clipping,
optional auxiliary metric trackers.

The overall objective optimized by this step is:

(19)#\[L_{total} = L_{data} + L_{phys}\]

where \(L_{data}\) is the compiled supervised loss and \(L_{phys}\) is the scaled physics loss returned by physics_core().

Parameters:

data (tuple) –

Keras batch payload as (inputs, targets).

inputs is a dict of tensors matching the GeoPrior input API (static, dynamic, future, coords, thickness, etc.).
targets is a dict (or dict-like) of supervised targets.

The method expects a dict-style multi-output target structure. Targets are canonicalized and reordered to match self.output_names.

Returns:

metrics – Dictionary of scalar tensors suitable for Keras logging. The exact keys are produced by pack_step_results() and typically include:

loss / total_loss: total objective value.
per-output supervised losses and metrics (from self.compiled_loss and self.compiled_metrics).
physics summary terms (e.g., physics_loss_scaled and selected components) when physics is enabled.
optional “manual” metrics from add-on trackers.

Return type:

Notes

Step outline. This training step performs the following stages:

Unpack and canonicalize targets
Targets are normalized into a stable dict structure using _canonicalize_targets and reordered by self._order_by_output_keys. Only keys in self.output_names are retained to guarantee consistent ordering for both loss computation and logging.
Forward pass with physics precomputation
The step calls physics_core() inside a single outer GradientTape. The physics core performs its own inner tape to compute coordinate derivatives required by PDE residuals. The outer tape ensures gradients flow through both:
- supervised data predictions, and
- physics loss scalars produced by the physics pathway.
Supervised data loss
Targets are aligned to prediction shapes (including quantile layout when applicable) using _align_true_for_loss and then passed as lists to self.compiled_loss. This allows Keras to apply:
- per-output losses configured in compile,
- regularization losses in self.losses,
- sample weighting logic if configured.
Total objective
The physics loss contribution is taken from the physics bundle as physics_loss_scaled. If physics is disabled (or gated off) the contribution is treated as zero.
Gradients, scaling, and clipping
Gradients of the total objective are computed w.r.t. all trainable variables. The step then:
- applies optional parameter-specific gradient scaling via self._scale_param_grads (for example, to slow down m_v or kappa updates),
- filters NaN/Inf gradients using filter_nan_gradients,
- applies global norm clipping (default clip value is 1.0),
- applies gradients via self.optimizer.apply_gradients.
This sequence is intended to improve stability for stiff physics losses and mixed-scale parameters.
Auxiliary trackers
If the model is configured with add-on trackers (for example, quantile coverage/sharpness or other custom diagnostics), update_state is called on the supervised outputs.
Packed return
The step returns a single packed dictionary from pack_step_results() so both training logs and evaluation summaries remain consistent.

Physics loss semantics. The physics contribution returned by physics_core() is already assembled with internal multipliers and (optionally) warmup/ramp gating. In other words, physics_loss_scaled is the quantity that should be added to the supervised loss.

If you need raw components for debugging, enable physics debug options in scaling_kwargs (for example, debug_physics_grads=True) and use the debug hooks called inside this step.

Gradient sanity and debugging. This method provides multiple stability and debug mechanisms:

NaN/Inf gradient filtering before applying updates.
Global-norm clipping to limit catastrophic updates.
Optional per-term gradient checks via dbg_term_grads_finite when scaling_kwargs['debug_physics_grads'] is enabled.

These are particularly useful when PDE residuals are large early in training or when coordinate scaling is misconfigured.

Examples

Typical usage: compile and fit normally, relying on this custom train step:

>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(1e-3),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
... )
>>> history = model.fit(train_ds, validation_data=val_ds, epochs=5)

Inspect returned metrics keys during training:

>>> logs = model.train_step(next(iter(train_ds)))
>>> sorted(list(logs))[:5]
['data_loss', 'loss', 'physics_loss_scaled', 'total_loss', ...]

See also

geoprior.models.subsidence.step_core.physics_core: Shared physics pathway used to compute PDE residuals and physics loss scalars consistently across train and eval.
pack_step_results: Pack supervised metrics, physics terms, and manual trackers into a stable Keras logging dictionary.
filter_nan_gradients: Sanitize gradient lists by removing NaN/Inf tensors.
tf.clip_by_global_norm: TensorFlow utility for global-norm gradient clipping.

test_step(data)[source]

Run one evaluation (validation/test) step for GeoPrior models.

This method overrides the standard tf.keras.Model.test_step to evaluate GeoPrior-style PINN models with dict inputs and multi-output targets. It computes:

supervised validation loss and metrics via compiled_loss and compiled metrics,
optional physics diagnostics and physics loss via _evaluate_physics_on_batch (no optimizer updates),
optional add-on tracker metrics (for example, quantile coverage and sharpness),
a unified packed logging dictionary returned by pack_step_results().

Unlike train_step(), this method does not apply gradients or update model parameters. It may still use a GradientTape internally for physics derivatives when physics is enabled, but no optimizer step occurs.

Parameters:

data (tuple) –

Keras batch payload as (inputs, targets).

inputs is a dict of tensors matching the GeoPrior input API (static, dynamic, future, coords, thickness, etc.).
targets is a dict (or dict-like) of supervised targets.

Targets are canonicalized and reordered to match self.output_names for stable loss computation.

Returns:

metrics – Dictionary of scalar tensors suitable for Keras validation logging. The exact keys depend on configured losses, metrics, and physics settings, and are produced by pack_step_results().

Typical keys include:

loss / total_loss: total evaluation objective.
data_loss: supervised loss only.
per-output losses/metrics from Keras compiled configuration.
physics summary terms (for example physics_loss_scaled, epsilons) if physics is enabled.
custom tracker metrics if add-on trackers are enabled.

Return type:

Notes

Step outline. This evaluation step follows a stable, dict-safe flow:

Unpack and canonicalize targets
Targets are normalized into a stable dict structure and reordered by output key contract.
Forward pass (supervised only)
The method calls call() via self(inputs, training=False) to obtain supervised predictions only. Aux tensors are not returned here by design.
Supervised loss and metrics
Targets are aligned to prediction shapes using _align_true_for_loss and passed to compiled_loss as ordered lists to ensure consistent behavior across Keras versions and dict wrappers.
Add-on trackers (optional)
If configured, add-on trackers are updated with targets and predictions. These trackers are purely diagnostic and do not affect loss values unless explicitly integrated elsewhere.
Physics diagnostics (optional)
If physics is enabled, the method calls _evaluate_physics_on_batch(inputs, return_maps=False) to compute physics residual summaries and a scaled physics loss.

The total evaluation objective is then:

(20)#\[L_{total} = L_{data} + L_{phys}\]

where \(L_{phys}\) is the physics loss scalar returned by the physics evaluator.

The physics evaluator may use internal autodiff to compute PDE derivatives for residual diagnostics, but gradients are not used to update parameters in test_step.
Packed return
The method returns a single packed dictionary from pack_step_results() to keep training and validation logs consistent.

When to use physics in validation. Enabling physics during validation is useful to monitor:

PDE residual RMS values (epsilon metrics),
consistency priors (for example, time-scale prior),
bounds penalties and stability signals.

If validation speed is a concern, physics can be disabled with the model physics switch (for example, _physics_off() returning True), in which case only supervised losses/metrics are computed.

Examples

Standard evaluation with physics enabled:

>>> logs = model.test_step(next(iter(val_ds)))
>>> float(logs["data_loss"])
1.23
>>> float(logs["physics_loss_scaled"])
0.01

Disable physics for faster validation (model-specific switch):

>>> model._physics_off = lambda: True
>>> logs = model.test_step(next(iter(val_ds)))
>>> "physics_loss_scaled" in logs
False  # depends on pack_step_results configuration

See also

train_step: Custom training step that computes physics loss and applies gradients.
_evaluate_physics_on_batch: Evaluation-only physics routine that computes residual diagnostics without applying optimizer updates.
pack_step_results: Pack supervised metrics, physics terms, and manual trackers into a stable Keras logging dictionary.

evaluate_physics(inputs, return_maps=False, max_batches=None, batch_size=None)[source]

Evaluate physics diagnostics over a batch or a dataset.

This method computes physics-only diagnostics for GeoPrior-style PINN models. Supported input modes are:

a tf.data.Dataset whose scalar diagnostics are aggregated across batches;
a mapping of tensors or numpy-like arrays, optionally batched via batch_size;
a single pre-batched mapping that is evaluated once.

The returned values are intended for monitoring PDE consistency, prior adherence, and stability during training and validation.

Parameters:

inputs (dict or Dataset) –
Input payload used for physics evaluation.
- If a dict, it should follow the GeoPrior batch API and contain tensors, or array-like values when batch_size is provided.
- If a Dataset, each element should yield either an input dict or a tuple/list whose first element is the input dict.
return_maps (bool, default False) –
If True, include residual maps and learned field tensors.

In Dataset mode, maps are not aggregated across batches. The method returns maps from the last processed batch only to keep memory usage bounded and avoid ambiguous aggregation semantics.
max_batches (int or None, default None) –
Maximum number of dataset batches to process. If None, iterate through the entire dataset.

This option is useful for quick diagnostics on large datasets.
batch_size (int or None, default None) – If provided and inputs is a mapping of numpy-like arrays, wrap into a dataset and batch by this size before evaluation.

Returns:

out – Dictionary of physics diagnostics. In Dataset mode, scalar keys whose names start with 'loss_' or 'epsilon_' are aggregated by mean across processed batches. Example aggregated outputs include loss_cons, loss_gw, loss_prior, loss_smooth, loss_bounds, loss_mv, loss_q_reg, epsilon_cons, epsilon_gw, and epsilon_prior.

When return_maps=True, the output may also include maps from the last processed batch, such as residuals R_prior, R_cons, R_gw; learned fields K, Ss, tau; closure-prior fields tau_prior / tau_closure; and thickness fields H_field / H plus drainage thickness Hd. Map availability depends on the underlying physics computation and whether the batch contains the required inputs.

Return type:

dict of str to Tensor

Raises:

ValueError – If the underlying physics computation requires missing inputs (for example, thickness) or inputs have incompatible shapes.

Notes

Use this method to evaluate physics consistency independently of the supervised data loss. Typical use cases include monitoring residual RMS values, diagnosing unit or coordinate mismatches, validating bounds and priors, and generating physics maps for inspection.

This method does not compute supervised metrics. In Dataset mode, only scalar keys with loss_ or epsilon_ prefixes are aggregated across batches. Residual maps and learned fields are not aggregated; when return_maps=True, the method returns the maps from the last processed batch.

Examples

Evaluate physics scalars over a validation dataset:

>>> phys = model.evaluate_physics(val_ds, max_batches=10)
>>> float(phys["epsilon_prior"])
0.01

Evaluate physics and retrieve last-batch maps:

>>> phys = model.evaluate_physics(val_ds, return_maps=True, max_batches=1)
>>> phys["R_gw"].shape
TensorShape([B, H, 1])

Evaluate a single batch dictionary:

>>> phys = model.evaluate_physics(batch_dict, return_maps=False)
>>> sorted([k for k in phys if k.startswith("loss_")])[:3]
['loss_bounds', 'loss_cons', 'loss_gw']

Wrap numpy-like arrays into batches (mapping mode):

>>> phys = model.evaluate_physics(inputs_np, batch_size=256, max_batches=5)

See also

_evaluate_physics_on_batch: Per-batch physics diagnostics wrapper.
geoprior.models.subsidence.step_core.physics_core: Shared physics computation used for diagnostics and training.

current_mv()[source]

Return the current value of the compressibility \(m_v\).

This is a thin convenience wrapper around _mv_value(), which handles both the trainable (log-parameterized) and fixed-scalar cases.

Returns:: Scalar tensor representing \(m_v\) in linear space.
Return type:: tf.Tensor

current_kappa()[source]

Return the current value of the consistency coefficient \(\kappa\).

This is a thin convenience wrapper around _kappa_value(), which handles both the trainable (log-parameterized) and fixed-scalar cases.

Returns:: Scalar tensor representing \(\kappa\) in linear space.
Return type:: tf.Tensor

get_last_physics_fields()[source]: Returns the most recent physics fields and H used by the model call. Shapes: (B, H, 1) each, matching the last forward pass.

split_data_predictions(data_tensor)[source]

Split a combined supervised output tensor into subsidence and GWL components.

GeoPrior models often compute a single “data head” tensor whose last dimension concatenates multiple supervised targets:

(21)#\[y = [s, g]\]

where \(s\) is subsidence and \(g\) is groundwater level (or a GWL-like driver). This helper slices the last axis into:

subsidence prediction tensor s_pred
groundwater-level prediction tensor gwl_pred

The slicing is controlled by the model attributes self.output_subsidence_dim and self.output_gwl_dim.

Parameters:

data_tensor (Tensor) –

Combined supervised output tensor with last axis size output_subsidence_dim + output_gwl_dim.

Typical shapes include:

(B, H, D) for point predictions, where D = subs_dim + gwl_dim.
(B, H, Q, D) for quantile predictions. In this case, the slicing is still applied on the last dimension D.

Returns:

s_pred (Tensor) – Subsidence slice from data_tensor[..., :output_subsidence_dim].
gwl_pred (Tensor) – GWL slice from data_tensor[..., output_subsidence_dim:].

Return type:

tuple[Tensor, Tensor]

Notes

This method performs a pure tensor slice and does not apply any unit conversions. Unit handling is managed by scaling helpers elsewhere.
If quantiles are present, the Q axis is preserved and only the last axis is split.

Examples

Point outputs:

>>> y = tf.zeros([8, 3, 2])  # subs_dim=1, gwl_dim=1
>>> s_pred, gwl_pred = model.split_data_predictions(y)
>>> s_pred.shape, gwl_pred.shape
(TensorShape([8, 3, 1]), TensorShape([8, 3, 1]))

Quantile outputs:

>>> yq = tf.zeros([8, 3, 3, 2])  # (B,H,Q,D)
>>> s_pred, gwl_pred = model.split_data_predictions(yq)
>>> s_pred.shape, gwl_pred.shape
(TensorShape([8, 3, 3, 1]), TensorShape([8, 3, 3, 1]))

See also

split_physics_predictions: Split the physics-head tensor into (K, Ss, dlogtau, Q) logits.

split_physics_predictions(phys_means_raw_tensor)[source]

Split the combined physics-head tensor into per-field logits.

GeoPrior models predict a compact “physics head” tensor whose last dimension concatenates the raw logits for multiple physics fields. This helper slices that tensor into:

K_logits : hydraulic conductivity logits
Ss_logits : specific storage logits
dlogtau_logits : relaxation time offset logits
Q_logits : optional forcing / source-term logits

The canonical ordering is:

(22)#\[p = [K, S_s, dlogtau, Q]\]

where each component is typically 1-dimensional, i.e. shape (B, H, 1) per component.

Parameters:

phys_means_raw_tensor (Tensor) –

Combined physics-head tensor. Expected shape is typically:

(B, H, P) where P is the total physics output dimension.
Some callers may supply tensors with additional axes, but the slicing always occurs along the last axis.

Returns:

K_logits (Tensor) – Slice corresponding to the conductivity logits. Shape is (..., output_K_dim) and usually (B, H, 1).
Ss_logits (Tensor) – Slice corresponding to the storage logits. Shape is (..., output_Ss_dim) and usually (B, H, 1).
dlogtau_logits (Tensor) – Slice corresponding to the relaxation-time offset logits. Shape is (..., output_tau_dim) and usually (B, H, 1).
Q_logits (Tensor) – Slice corresponding to the forcing/source logits. Shape is (..., output_Q_dim) and usually (B, H, 1).

If Q is disabled or missing from the input tensor, a zeros tensor with the appropriate broadcastable shape is returned.

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

Notes

Backward compatibility and “always return Q”. This helper is designed so downstream physics code never needs to branch on whether Q exists.

If self.output_Q_dim <= 0, Q is treated as disabled and a zeros tensor shaped like K_logits[..., :1] is returned.
If Q is enabled but phys_means_raw_tensor does not contain enough channels to include Q (older checkpoints), Q is returned as zeros with the correct shape.

This allows PDE residual code to accept a consistent signature regardless of whether Q is actually trained.

Shape and dimension conventions. The slice widths are controlled by model attributes:

output_K_dim
output_Ss_dim
output_tau_dim
output_Q_dim (optional)

If your model uses multi-dimensional physics heads, the returned tensors will preserve those widths accordingly.

Examples

Standard case with Q present:

>>> p = tf.zeros([8, 3, 4])  # [K,Ss,dlogtau,Q]
>>> K, Ss, dlogtau, Q = model.split_physics_predictions(p)
>>> K.shape, Ss.shape, dlogtau.shape, Q.shape
(TensorShape([8, 3, 1]), TensorShape([8, 3, 1]),
 TensorShape([8, 3, 1]), TensorShape([8, 3, 1]))

Backward-compatible case (no Q channel in stored tensor):

>>> p_old = tf.zeros([8, 3, 3])  # [K,Ss,dlogtau]
>>> K, Ss, dlogtau, Q = model.split_physics_predictions(p_old)
>>> Q.shape
TensorShape([8, 3, 1])

See also

compose_physics_fields: Map raw logits into bounded SI-consistent physics fields.
q_to_gw_source_term_si: Convert Q logits to the SI source term used in the GW PDE.

property lambda_offset_value: float: Current raw value stored in the TF weight _lambda_offset.

property lambda_offset: float

help(**kwargs)

property mv_lr_mult: float

Learning-rate multiplier for \(m_v\) (via log_mv).

This factor multiplies the gradient of the log-parameter log_mv inside _scale_param_grads(), allowing \(m_v\) to learn faster or slower than the rest of the network.

Returns:: Current value of the multiplier for log_mv.
Return type:: float

my_params = GeoPriorSubsNet( static_input_dim, dynamic_input_dim, future_input_dim, output_subsidence_dim=1, output_gwl_dim=1, embed_dim=32, hidden_units=64, lstm_units=64, attention_units=32, num_heads=4, dropout_rate=0.1, forecast_horizon=1, quantiles=None, max_window_size=10, memory_size=100, scales=None, multi_scale_agg='last', final_agg='last', activation='relu', use_residuals=True, use_batch_norm=False, pde_mode='both', identifiability_regime=None, mv=LearnableMV(initial_value=1e-07, trainable=True, name=learnable_mv), kappa=LearnableKappa(initial_value=1.0, trainable=True, name=learnable_kappa), gamma_w=FixedGammaW(value=9810.0, name=fixed_gamma_w, log_transform=True, non_negative=True), h_ref=FixedHRef(value=0.0, name=fixed_h_ref, log_transform=False, non_negative=False), use_effective_h=False, hd_factor=1.0, kappa_mode='kb', offset_mode='mul', bounds_mode='soft', residual_method='exact', time_units=None, use_vsn=True, vsn_units=None, mode=None, objective=None, attention_levels=None, architecture_config=None, scale_pde_residuals=True, scaling_kwargs=None, name='GeoPriorSubsNet', verbose=0 )

property kappa_lr_mult: float

Learning-rate multiplier for \(\kappa\) (via log_kappa).

This factor multiplies the gradient of the log-parameter log_kappa inside _scale_param_grads(), allowing \(\kappa\) to learn at a different pace than the other parameters.

Returns:: Current value of the multiplier for log_kappa.
Return type:: float

compile(lambda_cons=None, lambda_gw=None, lambda_prior=None, lambda_smooth=None, lambda_mv=None, lambda_bounds=None, lambda_q=None, lambda_offset=1.0, mv_lr_mult=1.0, kappa_lr_mult=1.0, scale_mv_with_offset=False, scale_q_with_offset=True, **kwargs)[source]

Compile the model and configure data/physics loss weighting.

This override extends tf.keras.Model.compile() with explicit weights for each physics term used by GeoPrior PINN training, plus a global physics multiplier (lambda_offset) that can be scheduled during training.

The GeoPrior training objective (as used by train_step()) is:

(23)#\[L_{total} = L_{data} + \alpha(\text{offset_mode}, \lambda_{offset}) \, L_{phys}\]

where the physics objective is assembled from multiple components:

(24)#\[\begin{split}L_{phys} = &&\lambda_{cons} L_{cons}\\ && + \lambda_{gw} L_{gw}\\ && + \lambda_{prior} L_{prior}\\ && + \lambda_{smooth} L_{smooth}\\ && + \lambda_{mv} L_{mv}\\ && + \lambda_{bounds} L_{bounds}\\ && + \lambda_{q} L_{q}\\\end{split}\]

Each component corresponds to a residual (or penalty) computed in the shared physics core and summarized as mean-square values. The global multiplier \(alpha\) is determined by self.offset_mode:

offset_mode='mul' : \(\alpha = \lambda_{offset}\)
offset_mode='log10': \(\alpha = 10^{\lambda_{offset}}\)

The value of lambda_offset is stored in a non-trainable scalar weight self._lambda_offset (created via add_weight), which makes it safe to update during training from callbacks.

Parameters:

lambda_cons (float, default 1.0) –
Weight for the consolidation residual loss \(L_{cons}\).

This term penalizes the (scaled) consolidation residual \(R_{cons}\) derived from the settlement relaxation update, and is typically computed as:

(25)\[L_{cons} = E[ R_{cons}^2 ]\]
lambda_gw (float, default 1.0) –
Weight for the groundwater-flow residual loss \(L_{gw}\).

This term penalizes the (scaled) groundwater PDE residual \(R_{gw}\) of the form:

(26)\[R_{gw} = S_s \, \partial_t h - \nabla \cdot (K \nabla h) - Q\]

and is typically computed as:

(27)\[L_{gw} = E[ R_{gw}^2 ]\]
lambda_prior (float, default 1.0) –
Weight for the consistency prior loss \(L_{prior}\).

This term ties the learned relaxation time \(tau\) to a closure-based timescale \(tau_{phys}\) computed from the learned fields and thickness. In the current implementation the residual is commonly expressed in log space:

(28)\[R_{prior} = \log(\tau) - \log(\tau_{phys})\]

and the loss is:

(29)\[L_{prior} = E[ R_{prior}^2 ]\]
lambda_smooth (float, default 1.0) –
Weight for the smoothness prior loss \(L_{smooth}\).

This term penalizes spatial roughness in the learned hydraulic fields, typically via squared first derivatives:

(30)\[L_{smooth} = E[ (\partial_x K)^2 + (\partial_y K)^2 + (\partial_x S_s)^2 + (\partial_y S_s)^2 ]\]

It stabilizes training and encourages spatially coherent fields.
lambda_mv (float, default 0.0) –
Weight for the m_v consistency prior \(L_{mv}\).

This term is designed to provide a direct learning signal for \(m_v\) by aligning \(S_s\) with the expected relation with compressibility and water unit weight:

(31)\[S_s \approx m_v \, \gamma_w\]

A common residual is constructed in log space for stability:

(32)\[R_{mv} = \log(S_s) - \log(m_v \gamma_w)\]

and the loss is:

(33)\[L_{mv} = E[ \rho(R_{mv}) ]\]

where \(rho\) may be a robust penalty (for example, Huber) depending on scaling_kwargs configuration. When set to a positive value, this term can help constrain \(m_v\) in underdetermined settings.
lambda_bounds (float, default 0.0) –
Weight for the bounds penalty \(L_{bounds}\).

This term penalizes violations of configured parameter bounds (for example, thickness and log-parameter ranges) provided in scaling_kwargs['bounds']. When bounds_mode='soft', the penalty is differentiable and contributes to the objective:

(34)\[L_{bounds} = E[ R_{bounds}^2 ]\]

When bounds_mode='hard', parameters may be clipped or projected by the physics mapping, and this weight is typically forced to zero.
lambda_q (float, default 0.0) –
Weight for the forcing regularization term \(L_{q}\).

This term discourages excessive forcing magnitude by penalizing the mean-square of the SI source term \(Q\) used in the GW residual:

(35)\[L_{q} = E[ Q^2 ]\]

It is useful when a learnable forcing head is enabled and you want it to remain near zero unless required by data.
lambda_offset (float, default 1.0) –
Global physics multiplier stored in self._lambda_offset.

The effective multiplier applied to \(L_{phys}\) is:
- for offset_mode='mul' : \(alpha = \lambda_{offset}\)
- for offset_mode='log10': \(alpha = 10^{\lambda_{offset}}\)
self._lambda_offset is a non-trainable scalar weight so it can be updated safely during training, for example:

model._lambda_offset.assign(new_value)
mv_lr_mult (float, default 1.0) – Learning-rate multiplier applied to the gradient updates of the m_v log-parameter. This affects only the parameter update scaling, not the loss definition.
kappa_lr_mult (float, default 1.0) – Learning-rate multiplier applied to the gradient updates of the kappa log-parameter (the closure/unit-conversion factor used by the timescale prior). This affects only parameter update scaling, not the loss definition.
scale_mv_with_offset (bool, default False) –
If True, multiply the \(L_{mv}\) contribution by the global physics multiplier \(alpha\) in addition to lambda_mv.

This is useful when \(L_{mv}\) should follow the same warmup schedule as other physics terms. If False, \(L_{mv}\) is weighted only by lambda_mv.
scale_q_with_offset (bool, default True) –
If True, multiply the \(L_{q}\) contribution by the global physics multiplier \(alpha\) in addition to lambda_q.

This is commonly enabled so forcing regularization ramps in together with other physics terms during physics warmup.
kwargs (dict) – Additional keyword arguments forwarded to tf.keras.Model.compile(), such as optimizer, loss, metrics, run_eagerly, jit_compile, and so on.

Returns:

self – Returns the compiled model instance.

Return type:

Notes

Physics-off behavior. If the model physics is disabled (for example, by PDE mode settings or a physics switch), this method forces all physics weights to neutral values regardless of the inputs:

lambda_prior = 0.0
lambda_smooth = 0.0
lambda_mv = 0.0
lambda_q = 0.0
lambda_bounds = 0.0
self._lambda_offset = 1.0

This ensures that train_step() and test_step() remain stable and that logs do not contain misleading physics terms.

Validation of lambda_offset. For offset_mode='mul', lambda_offset must be strictly positive. For offset_mode='log10', any real value is allowed and acts as a log10-scale controller.

Scheduling lambda_offset. A recommended pattern is to keep individual lambda_* values fixed and schedule lambda_offset (warmup/ramp) using a callback. Because self._lambda_offset is a non-trainable TF weight, it is safe to update at runtime.

Examples

Compile with physics enabled and a moderate prior:

>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(1e-3),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
...     lambda_cons=1.0,
...     lambda_gw=1.0,
...     lambda_prior=2.0,
...     lambda_smooth=0.1,
...     lambda_bounds=0.01,
...     lambda_offset=0.1,
... )

Disable forcing penalty and use a stronger smoothness prior:

>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(5e-4),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
...     lambda_q=0.0,
...     lambda_smooth=1.0,
... )

Use log10 scaling for the global physics multiplier:

>>> model.offset_mode = "log10"
>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(1e-3),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
...     lambda_offset=-1.0,  # physics multiplier = 0.1
... )

See also

train_step: Uses the configured lambdas to assemble the total loss and apply gradients.
_physics_loss_multiplier: Computes the global physics multiplier from offset_mode and self._lambda_offset.
geoprior.models.subsidence.step_core.physics_core: Computes per-batch physics residuals and loss terms.

export_physics_payload(dataset, max_batches=None, save_path=None, format='npz', overwrite=False, metadata=None, random_subsample=None, float_dtype=<class 'numpy.float32'>, log_fn=None, **tqdm_kws)[source]

Export physics diagnostics as a flat payload.

This helper collects physics diagnostics from a trained GeoPrior-style model and optionally persists them to disk.

Internally, it calls gather_physics_payload() to iterate over dataset and evaluate physics maps and scalar summaries via GeoPriorSubsNet.evaluate_physics() with return_maps=True. The per-batch tensors are flattened and concatenated into 1D arrays suitable for scatter plots, histograms, and reproducibility artifacts.

Parameters:

dataset (iterable) – Batched iterable (typically a tf.data.Dataset) yielding either inputs or (inputs, targets). Targets, if present, are ignored. Each inputs must contain the tensors required by evaluate_physics() (notably the coordinate tensor and thickness field, depending on the model configuration).
max_batches (int or None, default None) – Maximum number of batches to process. If None, consumes the entire iterable.
save_path (str or None, default None) – If provided, write the payload to this location using save_physics_payload(). If save_path is a directory, a default filename is used by the saver.
format ({'npz', 'csv', 'parquet'}, default 'npz') – Output format for persistence. 'npz' writes a compressed NumPy archive and a JSON sidecar metadata file.
overwrite (bool, default False) – If False and save_path already exists, raise an error.
metadata (dict or None, default None) – Optional user metadata to merge into the auto-generated provenance returned by default_meta_from_model(). User keys override defaults on conflict.
random_subsample (float or None, default None) – If provided, randomly subsample the flat payload after it is gathered. Must be in (0, 1] and is interpreted as the fraction of rows to keep. This is useful to reduce file size for large grids.
float_dtype (numpy dtype, default numpy.float32) – Dtype used when casting flattened arrays. Using float32 keeps files compact and is typically sufficient for diagnostics.
log_fn (callable or None, default None) – Optional logger used by the progress helper (for example, print). If None, the progress helper may be silent.
**tqdm_kws – Extra keyword arguments forwarded to the progress helper used inside gather_physics_payload().

Returns:

payload – Flat diagnostics payload with 1D arrays. The exact keys are defined by gather_physics_payload(), but typically include:

tau : effective relaxation time (seconds)
tau_prior / tau_closure : closure timescale (seconds)
K : effective hydraulic conductivity (m/s)
Ss : effective specific storage (1/m)
Hd : effective drainage thickness (m)
cons_res_vals : consolidation residual values
log10_tau and log10_tau_prior
metrics : nested dict with summary scalars

Return type:

dict[str, numpy.ndarray]

Notes

This routine does not change units. Unit consistency is a responsibility of the model physics and its scaling_kwargs.
If return_maps=True is used inside evaluate_physics(), maps are collected per batch and then flattened here. When saving, the payload is stored exactly as returned by the model.
Random subsampling is performed after concatenation, so it samples rows uniformly across all processed batches.

See also

gather_physics_payload: Core collector that builds the flat arrays.
save_physics_payload: Persist payload + metadata to disk.
default_meta_from_model: Build lightweight provenance metadata from a model.
GeoPriorSubsNet.evaluate_physics: Compute physics scalars and (optionally) maps.

Examples

>>> # ds is a batched tf.data.Dataset yielding (inputs, targets)
>>> payload = model.export_physics_payload(
...     ds, max_batches=20, random_subsample=0.25
... )
>>> # Save to disk (creates a .meta.json sidecar for npz/csv/parquet)
>>> _ = model.export_physics_payload(
...     ds,
...     max_batches=50,
...     save_path="physics_payload.npz",
...     format="npz",
...     overwrite=True,
... )

static load_physics_payload(path)[source]

Load a previously saved physics payload.

This is a thin convenience wrapper around load_physics_payload() from the diagnostics payload module. It reads the data file and its optional JSON sidecar metadata.

Parameters:

path (str) – Path to a saved payload. Supported extensions depend on the underlying loader and typically include .npz, .csv, and .parquet. For formats that support it, a sidecar metadata file is expected at path + '.meta.json'.

Returns:

(payload, meta) –

payloaddict[str, numpy.ndarray]: Dictionary of arrays loaded from disk. Backward- and forward-compatible aliases may be added by the loader (for example, ensuring both tau_prior and tau_closure are present).
metadict: Metadata dictionary loaded from the JSON sidecar if found, otherwise an empty dict.

Return type:

tuple(dict, dict)

Notes

This method performs I/O only. It does not validate that the payload matches a particular model instance.
If you saved with format='npz', the payload is loaded using NumPy. For CSV/Parquet, the loader typically uses pandas.

See also

load_physics_payload: The underlying loader that performs format dispatch.
GeoPriorSubsNet.export_physics_payload: Export and optionally save a payload.

Examples

>>> payload, meta = GeoPriorSubsNet.load_physics_payload(
...     "physics_payload.npz"
... )
>>> list(payload)[:5]
['tau', 'tau_prior', 'K', 'Ss', 'Hd']

get_config()[source]

Return a Keras-serializable configuration for model reconstruction.

This method extends tf.keras.Model.get_config() to ensure GeoPriorSubsNet can be saved and reloaded with tf.keras.models.load_model() (or keras.models.load_model()) while preserving the model’s physics options and scaling pipeline.

The returned dictionary contains:

the base configuration from BaseAttentive (via super().get_config()),
the supervised output layout (output_dim),
the resolved scaling configuration serialized as a Keras object,
GeoPrior-specific physics constructor arguments and flags.

The output is designed to be JSON-serializable by Keras. Objects that are not plain JSON (for example, GeoPriorScalingConfig and scalar wrappers such as LearnableMV) are included as Keras serialized objects via keras.saving.serialize_keras_object().

Returns:: config – A configuration dictionary that can be passed to from_config() to reconstruct the model.
Return type:: dict

Notes

output_dim is kept for compatibility with the BaseAttentive constructor signature. It is not a user-facing argument for the GeoPrior model; it is derived from:

(36)#\[output\_dim = output\_subsidence\_dim + output\_gwl\_dim\]
scaling_kwargs is stored as a serialized Keras object representing the validated scaling configuration. This preserves the exact conventions (units, coordinate normalization, bounds) used during training and is critical for consistent inference.
This config does not include runtime-only state such as optimizer variables or training metrics. Those are handled by standard Keras checkpointing mechanisms.

Examples

Serialize and reconstruct manually:

>>> cfg = model.get_config()
>>> model2 = model.__class__.from_config(cfg)

Save and reload through Keras:

>>> model.save("geoprior.keras")
>>> model2 = keras.models.load_model(
...     "geoprior.keras",
...     custom_objects={"GeoPriorSubsNet": GeoPriorSubsNet},
... )

See also

from_config: Reconstruct a model instance from the serialized config.
keras.saving.serialize_keras_object: Keras helper used to serialize non-JSON config objects.

classmethod from_config(config, custom_objects=None)[source]

Rebuild a GeoPrior model instance from a serialized configuration.

This classmethod reconstructs the model from a configuration dictionary produced by get_config() and used by the Keras serialization stack.

The method performs three reconstruction steps:

Build a custom_objects registry that includes all GeoPrior wrappers and scaling configuration classes needed for safe deserialization.
Rehydrate wrapper objects stored as Keras-serialized dicts ({"class_name": ..., "config": ...}) for keys such as mv, kappa, gamma_w, and h_ref.
Rehydrate the scaling configuration stored under scaling_kwargs if present as a Keras object.

Finally, the method removes legacy/internal keys that are not part of the current constructor signature and returns cls(**config).

Parameters:

config (dict) – Serialized configuration dictionary. Typically produced by get_config() and passed by Keras during deserialization.
custom_objects (dict or None, default None) – Optional mapping used by Keras to resolve custom layers, models, and config objects. If None, an internal registry is created and merged with any user-provided entries.

Returns:

model – A reconstructed model instance equivalent to the original model at save time (architecture and configuration). Weights are loaded by Keras separately when using keras.models.load_model().

Return type:

Notes

This method is designed to be robust to older saved configs by explicitly dropping keys that were used by previous GeoPrior/PINN variants (for example, legacy groundwater coefficient keys and internal version markers).
The deserialization process relies on Keras helpers and the custom_objects registry. If you have custom subclasses or external layers referenced inside architecture_config, you must provide them in custom_objects or register them with Keras before loading.
If scaling deserialization fails, the method raises the underlying exception because the scaling configuration is required for consistent unit handling and PDE residual computation.

Examples

Reconstruct from a saved config dictionary:

>>> cfg = model.get_config()
>>> model2 = GeoPriorSubsNet.from_config(
...     cfg,
...     custom_objects={"GeoPriorSubsNet": GeoPriorSubsNet},
... )

Load a saved model with explicit custom_objects:

>>> model2 = keras.models.load_model(
...     "geoprior.keras",
...     custom_objects={
...         "GeoPriorSubsNet": GeoPriorSubsNet,
...         "GeoPriorScalingConfig": GeoPriorScalingConfig,
...     },
... )

See also

get_config: Produce the configuration dictionary used for reconstruction.
keras.saving.deserialize_keras_object: Keras helper used to rehydrate serialized config objects.

class geoprior.models.subsidence.models.PoroElasticSubsNet(*args, **kwargs)[source]

Bases: GeoPriorSubsNet

Poroelastic surrogate variant of GeoPriorSubsNet.

This model is architecturally identical to GeoPriorSubsNet and follows the same dict-input API, outputs, and parameter semantics. It is provided as a physics-driven baseline for ablation and comparison runs.

Parameters:

static_input_dim (int)
dynamic_input_dim (int)
future_input_dim (int)
pde_mode (str)
use_effective_h (bool)
hd_factor (float)
kappa_mode (str)
scale_pde_residuals (bool)
scaling_kwargs (dict[str, Any] | None)
name (str)

help(**kwargs)

my_params = PoroElasticSubsNet( static_input_dim, dynamic_input_dim, future_input_dim, pde_mode='consolidation', use_effective_h=True, hd_factor=0.6, kappa_mode='bar', scale_pde_residuals=True, scaling_kwargs=None, name='PoroElasticSubsNet' )

__init__(static_input_dim, dynamic_input_dim, future_input_dim, pde_mode='consolidation', use_effective_h=True, hd_factor=0.6, kappa_mode='bar', scale_pde_residuals=True, scaling_kwargs=None, name='PoroElasticSubsNet', **kwargs)[source]

Parameters:

static_input_dim (int)
dynamic_input_dim (int)
future_input_dim (int)
pde_mode (str)
use_effective_h (bool)
hd_factor (float)
kappa_mode (str)
scale_pde_residuals (bool)
scaling_kwargs (dict[str, Any] | None)
name (str)

compile(lambda_cons=1.0, lambda_gw=0.0, lambda_prior=5.0, lambda_smooth=1.0, lambda_mv=0.1, lambda_bounds=0.05, mv_lr_mult=0.5, kappa_lr_mult=0.5, **kwargs)[source]

Compile with stronger defaults for the geomechanical prior.

Compared to GeoPriorSubsNet, this variant:

sets lambda_gw=0.0 (no groundwater-flow residual),
increases lambda_prior and lambda_bounds so that \(tau\) is tightly tied to \(tau_phys\),
gives \(m_v\) and \(kappa\) a smaller LR multiplier so they move more conservatively.

Parameters:

lambda_cons (float)
lambda_gw (float)
lambda_prior (float)
lambda_smooth (float)
lambda_mv (float)
lambda_bounds (float)
mv_lr_mult (float)
kappa_lr_mult (float)

Key model classes#

class geoprior.models.subsidence.models.GeoPriorSubsNet(*args, **kwargs)[source]

Bases: BaseAttentive

Prior-regularized physics-informed network for multi-step subsidence forecasting with groundwater coupling.

GeoPriorSubsNet combines a BaseAttentive encoder-decoder with a set of physics losses that constrain the forecast to respect a simplified groundwater-flow equation and a consolidation closure. In addition, it learns spatially varying physics fields and regularizes them against geologically motivated priors.

Parameters:

static_input_dim (int)
dynamic_input_dim (int)
future_input_dim (int)
output_subsidence_dim (int)
output_gwl_dim (int)
embed_dim (int)
hidden_units (int)
lstm_units (int)
attention_units (int)
num_heads (int)
dropout_rate (float)
forecast_horizon (int)
quantiles (list[float] | None)
max_window_size (int)
memory_size (int)
scales (list[int] | None)
multi_scale_agg (str)
final_agg (str)
activation (str)
use_residuals (bool)
use_batch_norm (bool)
pde_mode (str | list[str])
identifiability_regime (str | None)
mv (LearnableMV | float)
kappa (LearnableKappa | float)
gamma_w (FixedGammaW | float)
h_ref (FixedHRef | float | str | None)
use_effective_h (bool)
hd_factor (float)
kappa_mode (str)
offset_mode (str)
bounds_mode (str)
residual_method (str)
time_units (str | None)
use_vsn (bool)
vsn_units (int | None)
mode (str | None)
objective (str | None)
attention_levels (str | list[str] | None)
architecture_config (dict | None)
scale_pde_residuals (bool)
scaling_kwargs (dict[str, Any] | None)
name (str)
verbose (int)

OUTPUT_KEYS = ('subs_pred', 'gwl_pred')

__init__(static_input_dim, dynamic_input_dim, future_input_dim, output_subsidence_dim=1, output_gwl_dim=1, embed_dim=32, hidden_units=64, lstm_units=64, attention_units=32, num_heads=4, dropout_rate=0.1, forecast_horizon=1, quantiles=None, max_window_size=10, memory_size=100, scales=None, multi_scale_agg='last', final_agg='last', activation='relu', use_residuals=True, use_batch_norm=False, pde_mode='both', identifiability_regime=None, mv=LearnableMV(initial_value=1e-07, trainable=True, name=learnable_mv), kappa=LearnableKappa(initial_value=1.0, trainable=True, name=learnable_kappa), gamma_w=FixedGammaW(value=9810.0, name=fixed_gamma_w, log_transform=True, non_negative=True), h_ref=FixedHRef(value=0.0, name=fixed_h_ref, log_transform=False, non_negative=False), use_effective_h=False, hd_factor=1.0, kappa_mode='kb', offset_mode='mul', bounds_mode='soft', residual_method='exact', time_units=None, use_vsn=True, vsn_units=None, mode=None, objective=None, attention_levels=None, architecture_config=None, scale_pde_residuals=True, scaling_kwargs=None, name='GeoPriorSubsNet', verbose=0, **kwargs)[source]

Parameters:

static_input_dim (int)
dynamic_input_dim (int)
future_input_dim (int)
output_subsidence_dim (int)
output_gwl_dim (int)
embed_dim (int)
hidden_units (int)
lstm_units (int)
attention_units (int)
num_heads (int)
dropout_rate (float)
forecast_horizon (int)
quantiles (list[float] | None)
max_window_size (int)
memory_size (int)
scales (list[int] | None)
multi_scale_agg (str)
final_agg (str)
activation (str)
use_residuals (bool)
use_batch_norm (bool)
pde_mode (str | list[str])
identifiability_regime (str | None)
mv (LearnableMV | float)
kappa (LearnableKappa | float)
gamma_w (FixedGammaW | float)
h_ref (FixedHRef | float | str | None)
use_effective_h (bool)
hd_factor (float)
kappa_mode (str)
offset_mode (str)
bounds_mode (str)
residual_method (str)
time_units (str | None)
use_vsn (bool)
vsn_units (int | None)
mode (str | None)
objective (str | None)
attention_levels (str | list[str] | None)
architecture_config (dict | None)
scale_pde_residuals (bool)
scaling_kwargs (dict[str, Any] | None)
name (str)
verbose (int)

build(input_shape)[source]

Build the model’s weights and sublayers.

Keras may call build() (e.g. via model.build() or model.summary()) before the first forward pass. For subclassed models, we must ensure all sublayers are actually built, otherwise Keras can mark the layer as built while internal state remains unbuilt.

Parameters:: input_shape (Any)
Return type:: None

property metrics: List of all metrics.

run_encoder_decoder_core(static_input, dynamic_input, future_input, coords_input, training)[source]

Run the shared encoder-decoder core for GeoPrior inputs.

This override keeps the coordinate tensor aligned with the learned sequence features that are later consumed by the physics stack.

Parameters:

static_input (Tensor)
dynamic_input (Tensor)
future_input (Tensor)
coords_input (Tensor)
training (bool)

Return type:

tuple[Tensor, Tensor]

forward_with_aux(inputs, training=False)[source]

Return predictions and auxiliary tensors for diagnostics.

This method is a thin, public wrapper around _forward_all() that exposes both:

y_pred: the supervised outputs (what call() returns),
aux: intermediate tensors useful for debugging, physics evaluation, and research diagnostics.

Unlike call(), this method is intended for inspection and tooling. It does not change Keras training behavior because it does not alter loss computation or variable updates; it simply returns additional tensors already produced by the internal forward path.

Parameters:

inputs (dict) –
Dict-input batch compatible with GeoPrior PINN models.

Typical entries include:
- static_features : Tensor, shape (B, S)
- dynamic_features : Tensor, shape (B, H, D)
- future_features : Tensor, shape (B, H, F)
- coords : Tensor, shape (B, H, 3) with last axis ordered as (t, x, y)
- H_field or soil_thickness : Tensor, thickness field broadcastable to (B, H, 1)
The exact required keys depend on the model configuration and Stage-1 export. This wrapper delegates all parsing and validation to _forward_all().
training (bool, default False) – Forward-pass training flag. When True, dropout, batch norm, and other training-time layers behave accordingly.

Returns:

y_pred (dict of str to Tensor) – Supervised predictions in the same format as call(). At minimum, keys include 'subs_pred' and 'gwl_pred'.
aux (dict of str to Tensor) – Auxiliary tensors for diagnostics. Typical keys include:
- data_final: final data head tensor used for supervised outputs (may include quantile axis).
- data_mean_raw: mean-path output before quantile modeling.
- phys_mean_raw: concatenated physics logits (K, Ss, dlogtau, optional Q).
- phys_features_raw_3d: physics feature tensor emitted by the shared encoder-decoder core.

Return type:

tuple[dict[str, Tensor], dict[str, Tensor]]

Notes

This method is recommended for:

debugging NaN/Inf propagation (by inspecting aux),
computing physics residuals outside train_step using the same forward tensors,
building evaluation utilities that need intermediate heads.

Examples

Run a forward pass and inspect physics logits:

>>> y_pred, aux = model.forward_with_aux(batch, training=False)
>>> aux["phys_mean_raw"].shape
TensorShape([B, H, 4])

See also

call: Standard Keras forward that returns supervised outputs only.
_forward_all: Internal forward routine that returns both predictions and auxiliary tensors.

call(inputs, training=False)[source]

Keras forward method returning supervised outputs only.

This method defines the standard inference and training forward behavior expected by tf.keras.Model. It returns only the supervised output dictionary that participates in Keras loss computation and metric updates.

Internally, call() delegates to _forward_all() and discards the auxiliary outputs to ensure a stable, minimal prediction contract.

Parameters:

inputs (dict) –
Dict-input batch compatible with GeoPrior PINN models.

Typical entries include:
- static_features : Tensor, shape (B, S)
- dynamic_features : Tensor, shape (B, H, D)
- future_features : Tensor, shape (B, H, F)
- coords : Tensor, shape (B, H, 3) with last axis ordered as (t, x, y)
- H_field or soil_thickness : Tensor, thickness field
All parsing, shape checks, and coordinate handling are performed by _forward_all().
training (bool, default False) – Forward-pass training flag. When True, training-time behavior (dropout, batch norm, etc.) is enabled.

Returns:

y_pred – Supervised prediction dictionary. Keys are ordered by the model output contract (for example, ('subs_pred', 'gwl_pred')). Each tensor is typically shaped:

without quantiles: (B, H, 1)
with quantiles: (B, H, Q, 1) or a model-defined quantile layout

Return type:

dict of str to Tensor

Notes

Auxiliary tensors such as physics logits and intermediate features are intentionally excluded from the return value. Use forward_with_aux() when diagnostics are required.

Examples

Standard inference call:

>>> y = model(batch, training=False)
>>> sorted(y.keys())
['gwl_pred', 'subs_pred']

See also

forward_with_aux: Forward wrapper returning both predictions and diagnostics.
_forward_all: Internal routine returning (y_pred, aux).

train_step(data)[source]

Run one custom training step for GeoPrior-style PINN training.

This method overrides the standard tf.keras.Model.train_step to train a hybrid, physics-informed model with dict inputs and multi-output supervision. The step integrates:

supervised data losses (from compile / compiled_loss),
physics losses computed by physics_core(),
optional gradient scaling for selected parameters,
robust gradient sanitization and global-norm clipping,
optional auxiliary metric trackers.

The overall objective optimized by this step is:

(37)#\[L_{total} = L_{data} + L_{phys}\]

where \(L_{data}\) is the compiled supervised loss and \(L_{phys}\) is the scaled physics loss returned by physics_core().

Parameters:

data (tuple) –

Keras batch payload as (inputs, targets).

inputs is a dict of tensors matching the GeoPrior input API (static, dynamic, future, coords, thickness, etc.).
targets is a dict (or dict-like) of supervised targets.

The method expects a dict-style multi-output target structure. Targets are canonicalized and reordered to match self.output_names.

Returns:

metrics – Dictionary of scalar tensors suitable for Keras logging. The exact keys are produced by pack_step_results() and typically include:

loss / total_loss: total objective value.
per-output supervised losses and metrics (from self.compiled_loss and self.compiled_metrics).
physics summary terms (e.g., physics_loss_scaled and selected components) when physics is enabled.
optional “manual” metrics from add-on trackers.

Return type:

Notes

Step outline. This training step performs the following stages:

Unpack and canonicalize targets
Targets are normalized into a stable dict structure using _canonicalize_targets and reordered by self._order_by_output_keys. Only keys in self.output_names are retained to guarantee consistent ordering for both loss computation and logging.
Forward pass with physics precomputation
The step calls physics_core() inside a single outer GradientTape. The physics core performs its own inner tape to compute coordinate derivatives required by PDE residuals. The outer tape ensures gradients flow through both:
- supervised data predictions, and
- physics loss scalars produced by the physics pathway.
Supervised data loss
Targets are aligned to prediction shapes (including quantile layout when applicable) using _align_true_for_loss and then passed as lists to self.compiled_loss. This allows Keras to apply:
- per-output losses configured in compile,
- regularization losses in self.losses,
- sample weighting logic if configured.
Total objective
The physics loss contribution is taken from the physics bundle as physics_loss_scaled. If physics is disabled (or gated off) the contribution is treated as zero.
Gradients, scaling, and clipping
Gradients of the total objective are computed w.r.t. all trainable variables. The step then:
- applies optional parameter-specific gradient scaling via self._scale_param_grads (for example, to slow down m_v or kappa updates),
- filters NaN/Inf gradients using filter_nan_gradients,
- applies global norm clipping (default clip value is 1.0),
- applies gradients via self.optimizer.apply_gradients.
This sequence is intended to improve stability for stiff physics losses and mixed-scale parameters.
Auxiliary trackers
If the model is configured with add-on trackers (for example, quantile coverage/sharpness or other custom diagnostics), update_state is called on the supervised outputs.
Packed return
The step returns a single packed dictionary from pack_step_results() so both training logs and evaluation summaries remain consistent.

Physics loss semantics. The physics contribution returned by physics_core() is already assembled with internal multipliers and (optionally) warmup/ramp gating. In other words, physics_loss_scaled is the quantity that should be added to the supervised loss.

If you need raw components for debugging, enable physics debug options in scaling_kwargs (for example, debug_physics_grads=True) and use the debug hooks called inside this step.

Gradient sanity and debugging. This method provides multiple stability and debug mechanisms:

NaN/Inf gradient filtering before applying updates.
Global-norm clipping to limit catastrophic updates.
Optional per-term gradient checks via dbg_term_grads_finite when scaling_kwargs['debug_physics_grads'] is enabled.

These are particularly useful when PDE residuals are large early in training or when coordinate scaling is misconfigured.

Examples

Typical usage: compile and fit normally, relying on this custom train step:

>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(1e-3),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
... )
>>> history = model.fit(train_ds, validation_data=val_ds, epochs=5)

Inspect returned metrics keys during training:

>>> logs = model.train_step(next(iter(train_ds)))
>>> sorted(list(logs))[:5]
['data_loss', 'loss', 'physics_loss_scaled', 'total_loss', ...]

See also

geoprior.models.subsidence.step_core.physics_core: Shared physics pathway used to compute PDE residuals and physics loss scalars consistently across train and eval.
pack_step_results: Pack supervised metrics, physics terms, and manual trackers into a stable Keras logging dictionary.
filter_nan_gradients: Sanitize gradient lists by removing NaN/Inf tensors.
tf.clip_by_global_norm: TensorFlow utility for global-norm gradient clipping.

test_step(data)[source]

Run one evaluation (validation/test) step for GeoPrior models.

This method overrides the standard tf.keras.Model.test_step to evaluate GeoPrior-style PINN models with dict inputs and multi-output targets. It computes:

supervised validation loss and metrics via compiled_loss and compiled metrics,
optional physics diagnostics and physics loss via _evaluate_physics_on_batch (no optimizer updates),
optional add-on tracker metrics (for example, quantile coverage and sharpness),
a unified packed logging dictionary returned by pack_step_results().

Unlike train_step(), this method does not apply gradients or update model parameters. It may still use a GradientTape internally for physics derivatives when physics is enabled, but no optimizer step occurs.

Parameters:

data (tuple) –

Keras batch payload as (inputs, targets).

inputs is a dict of tensors matching the GeoPrior input API (static, dynamic, future, coords, thickness, etc.).
targets is a dict (or dict-like) of supervised targets.

Targets are canonicalized and reordered to match self.output_names for stable loss computation.

Returns:

metrics – Dictionary of scalar tensors suitable for Keras validation logging. The exact keys depend on configured losses, metrics, and physics settings, and are produced by pack_step_results().

Typical keys include:

loss / total_loss: total evaluation objective.
data_loss: supervised loss only.
per-output losses/metrics from Keras compiled configuration.
physics summary terms (for example physics_loss_scaled, epsilons) if physics is enabled.
custom tracker metrics if add-on trackers are enabled.

Return type:

Notes

Step outline. This evaluation step follows a stable, dict-safe flow:

Unpack and canonicalize targets
Targets are normalized into a stable dict structure and reordered by output key contract.
Forward pass (supervised only)
The method calls call() via self(inputs, training=False) to obtain supervised predictions only. Aux tensors are not returned here by design.
Supervised loss and metrics
Targets are aligned to prediction shapes using _align_true_for_loss and passed to compiled_loss as ordered lists to ensure consistent behavior across Keras versions and dict wrappers.
Add-on trackers (optional)
If configured, add-on trackers are updated with targets and predictions. These trackers are purely diagnostic and do not affect loss values unless explicitly integrated elsewhere.
Physics diagnostics (optional)
If physics is enabled, the method calls _evaluate_physics_on_batch(inputs, return_maps=False) to compute physics residual summaries and a scaled physics loss.

The total evaluation objective is then:

(38)#\[L_{total} = L_{data} + L_{phys}\]

where \(L_{phys}\) is the physics loss scalar returned by the physics evaluator.

The physics evaluator may use internal autodiff to compute PDE derivatives for residual diagnostics, but gradients are not used to update parameters in test_step.
Packed return
The method returns a single packed dictionary from pack_step_results() to keep training and validation logs consistent.

When to use physics in validation. Enabling physics during validation is useful to monitor:

PDE residual RMS values (epsilon metrics),
consistency priors (for example, time-scale prior),
bounds penalties and stability signals.

If validation speed is a concern, physics can be disabled with the model physics switch (for example, _physics_off() returning True), in which case only supervised losses/metrics are computed.

Examples

Standard evaluation with physics enabled:

>>> logs = model.test_step(next(iter(val_ds)))
>>> float(logs["data_loss"])
1.23
>>> float(logs["physics_loss_scaled"])
0.01

Disable physics for faster validation (model-specific switch):

>>> model._physics_off = lambda: True
>>> logs = model.test_step(next(iter(val_ds)))
>>> "physics_loss_scaled" in logs
False  # depends on pack_step_results configuration

See also

train_step: Custom training step that computes physics loss and applies gradients.
_evaluate_physics_on_batch: Evaluation-only physics routine that computes residual diagnostics without applying optimizer updates.
pack_step_results: Pack supervised metrics, physics terms, and manual trackers into a stable Keras logging dictionary.

evaluate_physics(inputs, return_maps=False, max_batches=None, batch_size=None)[source]

Evaluate physics diagnostics over a batch or a dataset.

This method computes physics-only diagnostics for GeoPrior-style PINN models. Supported input modes are:

a tf.data.Dataset whose scalar diagnostics are aggregated across batches;
a mapping of tensors or numpy-like arrays, optionally batched via batch_size;
a single pre-batched mapping that is evaluated once.

The returned values are intended for monitoring PDE consistency, prior adherence, and stability during training and validation.

Parameters:

inputs (dict or Dataset) –
Input payload used for physics evaluation.
- If a dict, it should follow the GeoPrior batch API and contain tensors, or array-like values when batch_size is provided.
- If a Dataset, each element should yield either an input dict or a tuple/list whose first element is the input dict.
return_maps (bool, default False) –
If True, include residual maps and learned field tensors.

In Dataset mode, maps are not aggregated across batches. The method returns maps from the last processed batch only to keep memory usage bounded and avoid ambiguous aggregation semantics.
max_batches (int or None, default None) –
Maximum number of dataset batches to process. If None, iterate through the entire dataset.

This option is useful for quick diagnostics on large datasets.
batch_size (int or None, default None) – If provided and inputs is a mapping of numpy-like arrays, wrap into a dataset and batch by this size before evaluation.

Returns:

out – Dictionary of physics diagnostics. In Dataset mode, scalar keys whose names start with 'loss_' or 'epsilon_' are aggregated by mean across processed batches. Example aggregated outputs include loss_cons, loss_gw, loss_prior, loss_smooth, loss_bounds, loss_mv, loss_q_reg, epsilon_cons, epsilon_gw, and epsilon_prior.

When return_maps=True, the output may also include maps from the last processed batch, such as residuals R_prior, R_cons, R_gw; learned fields K, Ss, tau; closure-prior fields tau_prior / tau_closure; and thickness fields H_field / H plus drainage thickness Hd. Map availability depends on the underlying physics computation and whether the batch contains the required inputs.

Return type:

dict of str to Tensor

Raises:

ValueError – If the underlying physics computation requires missing inputs (for example, thickness) or inputs have incompatible shapes.

Notes

Use this method to evaluate physics consistency independently of the supervised data loss. Typical use cases include monitoring residual RMS values, diagnosing unit or coordinate mismatches, validating bounds and priors, and generating physics maps for inspection.

This method does not compute supervised metrics. In Dataset mode, only scalar keys with loss_ or epsilon_ prefixes are aggregated across batches. Residual maps and learned fields are not aggregated; when return_maps=True, the method returns the maps from the last processed batch.

Examples

Evaluate physics scalars over a validation dataset:

>>> phys = model.evaluate_physics(val_ds, max_batches=10)
>>> float(phys["epsilon_prior"])
0.01

Evaluate physics and retrieve last-batch maps:

>>> phys = model.evaluate_physics(val_ds, return_maps=True, max_batches=1)
>>> phys["R_gw"].shape
TensorShape([B, H, 1])

Evaluate a single batch dictionary:

>>> phys = model.evaluate_physics(batch_dict, return_maps=False)
>>> sorted([k for k in phys if k.startswith("loss_")])[:3]
['loss_bounds', 'loss_cons', 'loss_gw']

Wrap numpy-like arrays into batches (mapping mode):

>>> phys = model.evaluate_physics(inputs_np, batch_size=256, max_batches=5)

See also

_evaluate_physics_on_batch: Per-batch physics diagnostics wrapper.
geoprior.models.subsidence.step_core.physics_core: Shared physics computation used for diagnostics and training.

current_mv()[source]

Return the current value of the compressibility \(m_v\).

This is a thin convenience wrapper around _mv_value(), which handles both the trainable (log-parameterized) and fixed-scalar cases.

Returns:: Scalar tensor representing \(m_v\) in linear space.
Return type:: tf.Tensor

current_kappa()[source]

Return the current value of the consistency coefficient \(\kappa\).

This is a thin convenience wrapper around _kappa_value(), which handles both the trainable (log-parameterized) and fixed-scalar cases.

Returns:: Scalar tensor representing \(\kappa\) in linear space.
Return type:: tf.Tensor

get_last_physics_fields()[source]: Returns the most recent physics fields and H used by the model call. Shapes: (B, H, 1) each, matching the last forward pass.

split_data_predictions(data_tensor)[source]

Split a combined supervised output tensor into subsidence and GWL components.

GeoPrior models often compute a single “data head” tensor whose last dimension concatenates multiple supervised targets:

(39)#\[y = [s, g]\]

where \(s\) is subsidence and \(g\) is groundwater level (or a GWL-like driver). This helper slices the last axis into:

subsidence prediction tensor s_pred
groundwater-level prediction tensor gwl_pred

The slicing is controlled by the model attributes self.output_subsidence_dim and self.output_gwl_dim.

Parameters:

data_tensor (Tensor) –

Combined supervised output tensor with last axis size output_subsidence_dim + output_gwl_dim.

Typical shapes include:

(B, H, D) for point predictions, where D = subs_dim + gwl_dim.
(B, H, Q, D) for quantile predictions. In this case, the slicing is still applied on the last dimension D.

Returns:

s_pred (Tensor) – Subsidence slice from data_tensor[..., :output_subsidence_dim].
gwl_pred (Tensor) – GWL slice from data_tensor[..., output_subsidence_dim:].

Return type:

tuple[Tensor, Tensor]

Notes

This method performs a pure tensor slice and does not apply any unit conversions. Unit handling is managed by scaling helpers elsewhere.
If quantiles are present, the Q axis is preserved and only the last axis is split.

Examples

Point outputs:

>>> y = tf.zeros([8, 3, 2])  # subs_dim=1, gwl_dim=1
>>> s_pred, gwl_pred = model.split_data_predictions(y)
>>> s_pred.shape, gwl_pred.shape
(TensorShape([8, 3, 1]), TensorShape([8, 3, 1]))

Quantile outputs:

>>> yq = tf.zeros([8, 3, 3, 2])  # (B,H,Q,D)
>>> s_pred, gwl_pred = model.split_data_predictions(yq)
>>> s_pred.shape, gwl_pred.shape
(TensorShape([8, 3, 3, 1]), TensorShape([8, 3, 3, 1]))

See also

split_physics_predictions: Split the physics-head tensor into (K, Ss, dlogtau, Q) logits.

split_physics_predictions(phys_means_raw_tensor)[source]

Split the combined physics-head tensor into per-field logits.

GeoPrior models predict a compact “physics head” tensor whose last dimension concatenates the raw logits for multiple physics fields. This helper slices that tensor into:

K_logits : hydraulic conductivity logits
Ss_logits : specific storage logits
dlogtau_logits : relaxation time offset logits
Q_logits : optional forcing / source-term logits

The canonical ordering is:

(40)#\[p = [K, S_s, dlogtau, Q]\]

where each component is typically 1-dimensional, i.e. shape (B, H, 1) per component.

Parameters:

phys_means_raw_tensor (Tensor) –

Combined physics-head tensor. Expected shape is typically:

(B, H, P) where P is the total physics output dimension.
Some callers may supply tensors with additional axes, but the slicing always occurs along the last axis.

Returns:

K_logits (Tensor) – Slice corresponding to the conductivity logits. Shape is (..., output_K_dim) and usually (B, H, 1).
Ss_logits (Tensor) – Slice corresponding to the storage logits. Shape is (..., output_Ss_dim) and usually (B, H, 1).
dlogtau_logits (Tensor) – Slice corresponding to the relaxation-time offset logits. Shape is (..., output_tau_dim) and usually (B, H, 1).
Q_logits (Tensor) – Slice corresponding to the forcing/source logits. Shape is (..., output_Q_dim) and usually (B, H, 1).

If Q is disabled or missing from the input tensor, a zeros tensor with the appropriate broadcastable shape is returned.

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

Notes

Backward compatibility and “always return Q”. This helper is designed so downstream physics code never needs to branch on whether Q exists.

If self.output_Q_dim <= 0, Q is treated as disabled and a zeros tensor shaped like K_logits[..., :1] is returned.
If Q is enabled but phys_means_raw_tensor does not contain enough channels to include Q (older checkpoints), Q is returned as zeros with the correct shape.

This allows PDE residual code to accept a consistent signature regardless of whether Q is actually trained.

Shape and dimension conventions. The slice widths are controlled by model attributes:

output_K_dim
output_Ss_dim
output_tau_dim
output_Q_dim (optional)

If your model uses multi-dimensional physics heads, the returned tensors will preserve those widths accordingly.

Examples

Standard case with Q present:

>>> p = tf.zeros([8, 3, 4])  # [K,Ss,dlogtau,Q]
>>> K, Ss, dlogtau, Q = model.split_physics_predictions(p)
>>> K.shape, Ss.shape, dlogtau.shape, Q.shape
(TensorShape([8, 3, 1]), TensorShape([8, 3, 1]),
 TensorShape([8, 3, 1]), TensorShape([8, 3, 1]))

Backward-compatible case (no Q channel in stored tensor):

>>> p_old = tf.zeros([8, 3, 3])  # [K,Ss,dlogtau]
>>> K, Ss, dlogtau, Q = model.split_physics_predictions(p_old)
>>> Q.shape
TensorShape([8, 3, 1])

See also

compose_physics_fields: Map raw logits into bounded SI-consistent physics fields.
q_to_gw_source_term_si: Convert Q logits to the SI source term used in the GW PDE.

property lambda_offset_value: float: Current raw value stored in the TF weight _lambda_offset.

property lambda_offset: float

help(**kwargs)

property mv_lr_mult: float

Learning-rate multiplier for \(m_v\) (via log_mv).

This factor multiplies the gradient of the log-parameter log_mv inside _scale_param_grads(), allowing \(m_v\) to learn faster or slower than the rest of the network.

Returns:: Current value of the multiplier for log_mv.
Return type:: float

my_params = GeoPriorSubsNet( static_input_dim, dynamic_input_dim, future_input_dim, output_subsidence_dim=1, output_gwl_dim=1, embed_dim=32, hidden_units=64, lstm_units=64, attention_units=32, num_heads=4, dropout_rate=0.1, forecast_horizon=1, quantiles=None, max_window_size=10, memory_size=100, scales=None, multi_scale_agg='last', final_agg='last', activation='relu', use_residuals=True, use_batch_norm=False, pde_mode='both', identifiability_regime=None, mv=LearnableMV(initial_value=1e-07, trainable=True, name=learnable_mv), kappa=LearnableKappa(initial_value=1.0, trainable=True, name=learnable_kappa), gamma_w=FixedGammaW(value=9810.0, name=fixed_gamma_w, log_transform=True, non_negative=True), h_ref=FixedHRef(value=0.0, name=fixed_h_ref, log_transform=False, non_negative=False), use_effective_h=False, hd_factor=1.0, kappa_mode='kb', offset_mode='mul', bounds_mode='soft', residual_method='exact', time_units=None, use_vsn=True, vsn_units=None, mode=None, objective=None, attention_levels=None, architecture_config=None, scale_pde_residuals=True, scaling_kwargs=None, name='GeoPriorSubsNet', verbose=0 )

property kappa_lr_mult: float

Learning-rate multiplier for \(\kappa\) (via log_kappa).

This factor multiplies the gradient of the log-parameter log_kappa inside _scale_param_grads(), allowing \(\kappa\) to learn at a different pace than the other parameters.

Returns:: Current value of the multiplier for log_kappa.
Return type:: float

compile(lambda_cons=None, lambda_gw=None, lambda_prior=None, lambda_smooth=None, lambda_mv=None, lambda_bounds=None, lambda_q=None, lambda_offset=1.0, mv_lr_mult=1.0, kappa_lr_mult=1.0, scale_mv_with_offset=False, scale_q_with_offset=True, **kwargs)[source]

Compile the model and configure data/physics loss weighting.

This override extends tf.keras.Model.compile() with explicit weights for each physics term used by GeoPrior PINN training, plus a global physics multiplier (lambda_offset) that can be scheduled during training.

The GeoPrior training objective (as used by train_step()) is:

(41)#\[L_{total} = L_{data} + \alpha(\text{offset_mode}, \lambda_{offset}) \, L_{phys}\]

where the physics objective is assembled from multiple components:

(42)#\[\begin{split}L_{phys} = &&\lambda_{cons} L_{cons}\\ && + \lambda_{gw} L_{gw}\\ && + \lambda_{prior} L_{prior}\\ && + \lambda_{smooth} L_{smooth}\\ && + \lambda_{mv} L_{mv}\\ && + \lambda_{bounds} L_{bounds}\\ && + \lambda_{q} L_{q}\\\end{split}\]

Each component corresponds to a residual (or penalty) computed in the shared physics core and summarized as mean-square values. The global multiplier \(alpha\) is determined by self.offset_mode:

offset_mode='mul' : \(\alpha = \lambda_{offset}\)
offset_mode='log10': \(\alpha = 10^{\lambda_{offset}}\)

The value of lambda_offset is stored in a non-trainable scalar weight self._lambda_offset (created via add_weight), which makes it safe to update during training from callbacks.

Parameters:

lambda_cons (float, default 1.0) –
Weight for the consolidation residual loss \(L_{cons}\).

This term penalizes the (scaled) consolidation residual \(R_{cons}\) derived from the settlement relaxation update, and is typically computed as:

(43)\[L_{cons} = E[ R_{cons}^2 ]\]
lambda_gw (float, default 1.0) –
Weight for the groundwater-flow residual loss \(L_{gw}\).

This term penalizes the (scaled) groundwater PDE residual \(R_{gw}\) of the form:

(44)\[R_{gw} = S_s \, \partial_t h - \nabla \cdot (K \nabla h) - Q\]

and is typically computed as:

(45)\[L_{gw} = E[ R_{gw}^2 ]\]
lambda_prior (float, default 1.0) –
Weight for the consistency prior loss \(L_{prior}\).

This term ties the learned relaxation time \(tau\) to a closure-based timescale \(tau_{phys}\) computed from the learned fields and thickness. In the current implementation the residual is commonly expressed in log space:

(46)\[R_{prior} = \log(\tau) - \log(\tau_{phys})\]

and the loss is:

(47)\[L_{prior} = E[ R_{prior}^2 ]\]
lambda_smooth (float, default 1.0) –
Weight for the smoothness prior loss \(L_{smooth}\).

This term penalizes spatial roughness in the learned hydraulic fields, typically via squared first derivatives:

(48)\[L_{smooth} = E[ (\partial_x K)^2 + (\partial_y K)^2 + (\partial_x S_s)^2 + (\partial_y S_s)^2 ]\]

It stabilizes training and encourages spatially coherent fields.
lambda_mv (float, default 0.0) –
Weight for the m_v consistency prior \(L_{mv}\).

This term is designed to provide a direct learning signal for \(m_v\) by aligning \(S_s\) with the expected relation with compressibility and water unit weight:

(49)\[S_s \approx m_v \, \gamma_w\]

A common residual is constructed in log space for stability:

(50)\[R_{mv} = \log(S_s) - \log(m_v \gamma_w)\]

and the loss is:

(51)\[L_{mv} = E[ \rho(R_{mv}) ]\]

where \(rho\) may be a robust penalty (for example, Huber) depending on scaling_kwargs configuration. When set to a positive value, this term can help constrain \(m_v\) in underdetermined settings.
lambda_bounds (float, default 0.0) –
Weight for the bounds penalty \(L_{bounds}\).

This term penalizes violations of configured parameter bounds (for example, thickness and log-parameter ranges) provided in scaling_kwargs['bounds']. When bounds_mode='soft', the penalty is differentiable and contributes to the objective:

(52)\[L_{bounds} = E[ R_{bounds}^2 ]\]

When bounds_mode='hard', parameters may be clipped or projected by the physics mapping, and this weight is typically forced to zero.
lambda_q (float, default 0.0) –
Weight for the forcing regularization term \(L_{q}\).

This term discourages excessive forcing magnitude by penalizing the mean-square of the SI source term \(Q\) used in the GW residual:

(53)\[L_{q} = E[ Q^2 ]\]

It is useful when a learnable forcing head is enabled and you want it to remain near zero unless required by data.
lambda_offset (float, default 1.0) –
Global physics multiplier stored in self._lambda_offset.

The effective multiplier applied to \(L_{phys}\) is:
- for offset_mode='mul' : \(alpha = \lambda_{offset}\)
- for offset_mode='log10': \(alpha = 10^{\lambda_{offset}}\)
self._lambda_offset is a non-trainable scalar weight so it can be updated safely during training, for example:

model._lambda_offset.assign(new_value)
mv_lr_mult (float, default 1.0) – Learning-rate multiplier applied to the gradient updates of the m_v log-parameter. This affects only the parameter update scaling, not the loss definition.
kappa_lr_mult (float, default 1.0) – Learning-rate multiplier applied to the gradient updates of the kappa log-parameter (the closure/unit-conversion factor used by the timescale prior). This affects only parameter update scaling, not the loss definition.
scale_mv_with_offset (bool, default False) –
If True, multiply the \(L_{mv}\) contribution by the global physics multiplier \(alpha\) in addition to lambda_mv.

This is useful when \(L_{mv}\) should follow the same warmup schedule as other physics terms. If False, \(L_{mv}\) is weighted only by lambda_mv.
scale_q_with_offset (bool, default True) –
If True, multiply the \(L_{q}\) contribution by the global physics multiplier \(alpha\) in addition to lambda_q.

This is commonly enabled so forcing regularization ramps in together with other physics terms during physics warmup.
kwargs (dict) – Additional keyword arguments forwarded to tf.keras.Model.compile(), such as optimizer, loss, metrics, run_eagerly, jit_compile, and so on.

Returns:

self – Returns the compiled model instance.

Return type:

Notes

Physics-off behavior. If the model physics is disabled (for example, by PDE mode settings or a physics switch), this method forces all physics weights to neutral values regardless of the inputs:

lambda_prior = 0.0
lambda_smooth = 0.0
lambda_mv = 0.0
lambda_q = 0.0
lambda_bounds = 0.0
self._lambda_offset = 1.0

This ensures that train_step() and test_step() remain stable and that logs do not contain misleading physics terms.

Validation of lambda_offset. For offset_mode='mul', lambda_offset must be strictly positive. For offset_mode='log10', any real value is allowed and acts as a log10-scale controller.

Scheduling lambda_offset. A recommended pattern is to keep individual lambda_* values fixed and schedule lambda_offset (warmup/ramp) using a callback. Because self._lambda_offset is a non-trainable TF weight, it is safe to update at runtime.

Examples

Compile with physics enabled and a moderate prior:

>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(1e-3),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
...     lambda_cons=1.0,
...     lambda_gw=1.0,
...     lambda_prior=2.0,
...     lambda_smooth=0.1,
...     lambda_bounds=0.01,
...     lambda_offset=0.1,
... )

Disable forcing penalty and use a stronger smoothness prior:

>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(5e-4),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
...     lambda_q=0.0,
...     lambda_smooth=1.0,
... )

Use log10 scaling for the global physics multiplier:

>>> model.offset_mode = "log10"
>>> model.compile(
...     optimizer=tf.keras.optimizers.Adam(1e-3),
...     loss={"subs_pred": "mse", "gwl_pred": "mse"},
...     lambda_offset=-1.0,  # physics multiplier = 0.1
... )

See also

train_step: Uses the configured lambdas to assemble the total loss and apply gradients.
_physics_loss_multiplier: Computes the global physics multiplier from offset_mode and self._lambda_offset.
geoprior.models.subsidence.step_core.physics_core: Computes per-batch physics residuals and loss terms.

export_physics_payload(dataset, max_batches=None, save_path=None, format='npz', overwrite=False, metadata=None, random_subsample=None, float_dtype=<class 'numpy.float32'>, log_fn=None, **tqdm_kws)[source]

Export physics diagnostics as a flat payload.

This helper collects physics diagnostics from a trained GeoPrior-style model and optionally persists them to disk.

Internally, it calls gather_physics_payload() to iterate over dataset and evaluate physics maps and scalar summaries via GeoPriorSubsNet.evaluate_physics() with return_maps=True. The per-batch tensors are flattened and concatenated into 1D arrays suitable for scatter plots, histograms, and reproducibility artifacts.

Parameters:

dataset (iterable) – Batched iterable (typically a tf.data.Dataset) yielding either inputs or (inputs, targets). Targets, if present, are ignored. Each inputs must contain the tensors required by evaluate_physics() (notably the coordinate tensor and thickness field, depending on the model configuration).
max_batches (int or None, default None) – Maximum number of batches to process. If None, consumes the entire iterable.
save_path (str or None, default None) – If provided, write the payload to this location using save_physics_payload(). If save_path is a directory, a default filename is used by the saver.
format ({'npz', 'csv', 'parquet'}, default 'npz') – Output format for persistence. 'npz' writes a compressed NumPy archive and a JSON sidecar metadata file.
overwrite (bool, default False) – If False and save_path already exists, raise an error.
metadata (dict or None, default None) – Optional user metadata to merge into the auto-generated provenance returned by default_meta_from_model(). User keys override defaults on conflict.
random_subsample (float or None, default None) – If provided, randomly subsample the flat payload after it is gathered. Must be in (0, 1] and is interpreted as the fraction of rows to keep. This is useful to reduce file size for large grids.
float_dtype (numpy dtype, default numpy.float32) – Dtype used when casting flattened arrays. Using float32 keeps files compact and is typically sufficient for diagnostics.
log_fn (callable or None, default None) – Optional logger used by the progress helper (for example, print). If None, the progress helper may be silent.
**tqdm_kws – Extra keyword arguments forwarded to the progress helper used inside gather_physics_payload().

Returns:

payload – Flat diagnostics payload with 1D arrays. The exact keys are defined by gather_physics_payload(), but typically include:

tau : effective relaxation time (seconds)
tau_prior / tau_closure : closure timescale (seconds)
K : effective hydraulic conductivity (m/s)
Ss : effective specific storage (1/m)
Hd : effective drainage thickness (m)
cons_res_vals : consolidation residual values
log10_tau and log10_tau_prior
metrics : nested dict with summary scalars

Return type:

dict[str, numpy.ndarray]

Notes

This routine does not change units. Unit consistency is a responsibility of the model physics and its scaling_kwargs.
If return_maps=True is used inside evaluate_physics(), maps are collected per batch and then flattened here. When saving, the payload is stored exactly as returned by the model.
Random subsampling is performed after concatenation, so it samples rows uniformly across all processed batches.

See also

gather_physics_payload: Core collector that builds the flat arrays.
save_physics_payload: Persist payload + metadata to disk.
default_meta_from_model: Build lightweight provenance metadata from a model.
GeoPriorSubsNet.evaluate_physics: Compute physics scalars and (optionally) maps.

Examples

>>> # ds is a batched tf.data.Dataset yielding (inputs, targets)
>>> payload = model.export_physics_payload(
...     ds, max_batches=20, random_subsample=0.25
... )
>>> # Save to disk (creates a .meta.json sidecar for npz/csv/parquet)
>>> _ = model.export_physics_payload(
...     ds,
...     max_batches=50,
...     save_path="physics_payload.npz",
...     format="npz",
...     overwrite=True,
... )

static load_physics_payload(path)[source]

Load a previously saved physics payload.

This is a thin convenience wrapper around load_physics_payload() from the diagnostics payload module. It reads the data file and its optional JSON sidecar metadata.

Parameters:

path (str) – Path to a saved payload. Supported extensions depend on the underlying loader and typically include .npz, .csv, and .parquet. For formats that support it, a sidecar metadata file is expected at path + '.meta.json'.

Returns:

(payload, meta) –

payloaddict[str, numpy.ndarray]: Dictionary of arrays loaded from disk. Backward- and forward-compatible aliases may be added by the loader (for example, ensuring both tau_prior and tau_closure are present).
metadict: Metadata dictionary loaded from the JSON sidecar if found, otherwise an empty dict.

Return type:

tuple(dict, dict)

Notes

This method performs I/O only. It does not validate that the payload matches a particular model instance.
If you saved with format='npz', the payload is loaded using NumPy. For CSV/Parquet, the loader typically uses pandas.

See also

load_physics_payload: The underlying loader that performs format dispatch.
GeoPriorSubsNet.export_physics_payload: Export and optionally save a payload.

Examples

>>> payload, meta = GeoPriorSubsNet.load_physics_payload(
...     "physics_payload.npz"
... )
>>> list(payload)[:5]
['tau', 'tau_prior', 'K', 'Ss', 'Hd']

get_config()[source]

Return a Keras-serializable configuration for model reconstruction.

This method extends tf.keras.Model.get_config() to ensure GeoPriorSubsNet can be saved and reloaded with tf.keras.models.load_model() (or keras.models.load_model()) while preserving the model’s physics options and scaling pipeline.

The returned dictionary contains:

the base configuration from BaseAttentive (via super().get_config()),
the supervised output layout (output_dim),
the resolved scaling configuration serialized as a Keras object,
GeoPrior-specific physics constructor arguments and flags.

The output is designed to be JSON-serializable by Keras. Objects that are not plain JSON (for example, GeoPriorScalingConfig and scalar wrappers such as LearnableMV) are included as Keras serialized objects via keras.saving.serialize_keras_object().

Returns:: config – A configuration dictionary that can be passed to from_config() to reconstruct the model.
Return type:: dict

Notes

output_dim is kept for compatibility with the BaseAttentive constructor signature. It is not a user-facing argument for the GeoPrior model; it is derived from:

(54)#\[output\_dim = output\_subsidence\_dim + output\_gwl\_dim\]
scaling_kwargs is stored as a serialized Keras object representing the validated scaling configuration. This preserves the exact conventions (units, coordinate normalization, bounds) used during training and is critical for consistent inference.
This config does not include runtime-only state such as optimizer variables or training metrics. Those are handled by standard Keras checkpointing mechanisms.

Examples

Serialize and reconstruct manually:

>>> cfg = model.get_config()
>>> model2 = model.__class__.from_config(cfg)

Save and reload through Keras:

>>> model.save("geoprior.keras")
>>> model2 = keras.models.load_model(
...     "geoprior.keras",
...     custom_objects={"GeoPriorSubsNet": GeoPriorSubsNet},
... )

See also

from_config: Reconstruct a model instance from the serialized config.
keras.saving.serialize_keras_object: Keras helper used to serialize non-JSON config objects.

classmethod from_config(config, custom_objects=None)[source]

Rebuild a GeoPrior model instance from a serialized configuration.

This classmethod reconstructs the model from a configuration dictionary produced by get_config() and used by the Keras serialization stack.

The method performs three reconstruction steps:

Build a custom_objects registry that includes all GeoPrior wrappers and scaling configuration classes needed for safe deserialization.
Rehydrate wrapper objects stored as Keras-serialized dicts ({"class_name": ..., "config": ...}) for keys such as mv, kappa, gamma_w, and h_ref.
Rehydrate the scaling configuration stored under scaling_kwargs if present as a Keras object.

Finally, the method removes legacy/internal keys that are not part of the current constructor signature and returns cls(**config).

Parameters:

config (dict) – Serialized configuration dictionary. Typically produced by get_config() and passed by Keras during deserialization.
custom_objects (dict or None, default None) – Optional mapping used by Keras to resolve custom layers, models, and config objects. If None, an internal registry is created and merged with any user-provided entries.

Returns:

model – A reconstructed model instance equivalent to the original model at save time (architecture and configuration). Weights are loaded by Keras separately when using keras.models.load_model().

Return type:

Notes

This method is designed to be robust to older saved configs by explicitly dropping keys that were used by previous GeoPrior/PINN variants (for example, legacy groundwater coefficient keys and internal version markers).
The deserialization process relies on Keras helpers and the custom_objects registry. If you have custom subclasses or external layers referenced inside architecture_config, you must provide them in custom_objects or register them with Keras before loading.
If scaling deserialization fails, the method raises the underlying exception because the scaling configuration is required for consistent unit handling and PDE residual computation.

Examples

Reconstruct from a saved config dictionary:

>>> cfg = model.get_config()
>>> model2 = GeoPriorSubsNet.from_config(
...     cfg,
...     custom_objects={"GeoPriorSubsNet": GeoPriorSubsNet},
... )

Load a saved model with explicit custom_objects:

>>> model2 = keras.models.load_model(
...     "geoprior.keras",
...     custom_objects={
...         "GeoPriorSubsNet": GeoPriorSubsNet,
...         "GeoPriorScalingConfig": GeoPriorScalingConfig,
...     },
... )

See also

get_config: Produce the configuration dictionary used for reconstruction.
keras.saving.deserialize_keras_object: Keras helper used to rehydrate serialized config objects.

class geoprior.models.subsidence.models.PoroElasticSubsNet(*args, **kwargs)[source]

Bases: GeoPriorSubsNet

Poroelastic surrogate variant of GeoPriorSubsNet.

This model is architecturally identical to GeoPriorSubsNet and follows the same dict-input API, outputs, and parameter semantics. It is provided as a physics-driven baseline for ablation and comparison runs.

Parameters:

static_input_dim (int)
dynamic_input_dim (int)
future_input_dim (int)
pde_mode (str)
use_effective_h (bool)
hd_factor (float)
kappa_mode (str)
scale_pde_residuals (bool)
scaling_kwargs (dict[str, Any] | None)
name (str)

help(**kwargs)

my_params = PoroElasticSubsNet( static_input_dim, dynamic_input_dim, future_input_dim, pde_mode='consolidation', use_effective_h=True, hd_factor=0.6, kappa_mode='bar', scale_pde_residuals=True, scaling_kwargs=None, name='PoroElasticSubsNet' )

__init__(static_input_dim, dynamic_input_dim, future_input_dim, pde_mode='consolidation', use_effective_h=True, hd_factor=0.6, kappa_mode='bar', scale_pde_residuals=True, scaling_kwargs=None, name='PoroElasticSubsNet', **kwargs)[source]

Parameters:

static_input_dim (int)
dynamic_input_dim (int)
future_input_dim (int)
pde_mode (str)
use_effective_h (bool)
hd_factor (float)
kappa_mode (str)
scale_pde_residuals (bool)
scaling_kwargs (dict[str, Any] | None)
name (str)

compile(lambda_cons=1.0, lambda_gw=0.0, lambda_prior=5.0, lambda_smooth=1.0, lambda_mv=0.1, lambda_bounds=0.05, mv_lr_mult=0.5, kappa_lr_mult=0.5, **kwargs)[source]

Compile with stronger defaults for the geomechanical prior.

Compared to GeoPriorSubsNet, this variant:

sets lambda_gw=0.0 (no groundwater-flow residual),
increases lambda_prior and lambda_bounds so that \(tau\) is tightly tied to \(tau_phys\),
gives \(m_v\) and \(kappa\) a smaller LR multiplier so they move more conservatively.

Parameters:

lambda_cons (float)
lambda_gw (float)
lambda_prior (float)
lambda_smooth (float)
lambda_mv (float)
lambda_bounds (float)
mv_lr_mult (float)
kappa_lr_mult (float)

Scientific math helpers#

The maths module contains the low-level mathematical helpers used to:

compose effective physical fields,
derive closure timescales,
compute equilibrium compaction,
map forcing terms into groundwater-source form, and
compute soft-bounds residuals.

This layer is the most compact expression of the physical assumptions used by the model family, so it is often the best place to start when you want to understand how the learned fields are turned into physically meaningful quantities.

GeoPrior maths helpers (physics terms + scaling).

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:

(55)#\[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:

(56)#\[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:

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

The constraint is applied in log space for numerical stability:

(58)#\[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:

(59)#\[\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():

(60)#\[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:

(61)#\[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:

Positivity: \(\tau > 0\) is enforced implicitly.
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:

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

where, in 2D:

(63)#\[\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:

(64)#\[\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:

Explicit log bounds: log_tau_min and log_tau_max (already log-seconds).
Linear bounds in seconds: tau_min and tau_max.
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.
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.:

(65)#\[\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:

(66)#\[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:

x (Tensor)
eps (float)

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:

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

Define the raw log-parameters (logits) as:

(68)#\[\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:

(69)#\[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:

(70)#\[\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:

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

and the total learned timescale is:

(72)#\[\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:

(73)#\[\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:

(74)#\[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:

key (str)
scaling_kwargs (dict[str, Any] | None)
default_val (float | str)

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:

(75)\[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:

subsidence_kind="cumulative"

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

(76)#\[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)\).
subsidence_kind="increment"

s_pred_si is interpreted as per-step increments:

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

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

(78)#\[\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:

(79)#\[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:

(80)#\[\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:

(81)\[\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:

(82)\[\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:

(83)\[\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:
  
  (84)#\[\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:

x (Tensor)
eps (float)

Return type:

Tensor

Important helper functions#

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:

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

Define the raw log-parameters (logits) as:

(86)#\[\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:

(87)#\[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:

(88)#\[\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:

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

and the total learned timescale is:

(90)#\[\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:

(91)#\[\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.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.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.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:

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

The constraint is applied in log space for numerical stability:

(93)#\[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.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:

(94)#\[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.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]

Scaling and serialization#

The scaling layer defines the contract that tells GeoPrior what the data mean physically. It is one of the most important pieces of the whole subsidence stack, because it is not merely a preprocessing detail: it records the interpretation needed to connect model-space quantities to SI units, coordinates, bounds, and downstream diagnostics.

GeoPrior scaling config helpers (Keras-serializable).

class geoprior.models.subsidence.scaling.GeoPriorScalingConfig(payload=<factory>, source=None, schema_version='1')[source]

Bases: object

Scaling configuration utilities for GeoPrior PINN.

This module defines GeoPriorScalingConfig, a small Keras-serializable container used to store and reconstruct the physics scaling and slicing controls used by GeoPriorSubsNet.

The scaling configuration is critical because it governs how coordinates, time units, groundwater variables, and physics residuals are interpreted and non-dimensionalized. If this configuration is not faithfully serialized via Keras get_config(), a reloaded model may be reconstructed with a different effective physics behavior.

The main entry point is GeoPriorScalingConfig.from_any(), which accepts a dict-like mapping, a file path str, or an existing GeoPriorScalingConfig instance. The resolved configuration is produced by resolve(), which runs the same canonicalization and validation pipeline used during training.

Notes

The resolved scaling dictionary should be JSON-safe and stable under Keras serialization.
Use _jsonify() to defensively convert nested values (NumPy scalars, tuples, sets) into plain Python types.
The config container combines Keras serialization patterns with the standard-library dataclass model [16, 17].

See also

load_scaling_kwargs: Load scaling configuration from mapping or file.
canonicalize_scaling_kwargs: Normalize keys and fill defaults consistently.
enforce_scaling_alias_consistency: Ensure alias keys agree and do not conflict.
validate_scaling_kwargs: Validate schema and value ranges.

Parameters:

payload (dict)
source (str | None)
schema_version (str)

payload: dict

source: str | None = None

schema_version: str = '1'

classmethod from_any(obj, *, copy=True)[source]

Serializable container for GeoPrior scaling configuration.

This dataclass stores a “payload” dictionary that holds all scaling and physics-control parameters required to reproduce the model behavior after saving and reloading with Keras.

The container supports flexible construction from: - None (empty config), - a mapping (dict-like), - a file path str (loaded via load_scaling_kwargs), - an existing GeoPriorScalingConfig instance.

The canonical and validated configuration is produced by resolve(), which applies the GeoPrior scaling pipeline: loading, canonicalization, alias consistency checks, and validation.

Parameters:

payload (dict, optional) – Raw scaling configuration payload. This may be incomplete or contain aliases prior to canonicalization.
source (str or None, optional) – Optional provenance string, typically a file path used to load the payload. This is stored for traceability only.
schema_version (str, optional) – Version label for the payload schema. This can be used to implement migrations when the scaling format evolves.

Variables:

payload (dict) – The raw payload stored in this object.
source (str or None) – The provenance hint, if provided.
schema_version (str) – Schema version label.

Notes

The resolved scaling dictionary returned by resolve() is the one you should pass to the model internals.
get_config returns JSON-safe objects only. This avoids subtle reconstruction drift caused by non-serializable values.
This factory aligns with the Keras object-serialization pattern described in Keras Team [16].

Examples

Construct from a mapping:

>>> cfg = GeoPriorScalingConfig.from_any(
...     {"coords_normalized": True}
... )
>>> sk = cfg.resolve()
>>> isinstance(sk, dict)
True

Construct from a file path:

>>> cfg = GeoPriorScalingConfig.from_any(
...     "path/to/scaling_kwargs.json"
... )
>>> sk = cfg.resolve()

Use in a model constructor (pattern):

>>> cfg = GeoPriorScalingConfig.from_any(scaling_kwargs)
>>> scaling_kwargs_resolved = cfg.resolve()

See also

GeoPriorScalingConfig.from_any: Build config from dict, path, or config instance.
GeoPriorScalingConfig.resolve: Produce canonical and validated scaling dictionary.

load_scaling_kwargs, canonicalize_scaling_kwargs

resolve()[source]

Resolve the payload into a canonical, validated scaling dict.

This method runs the GeoPrior scaling pipeline and returns a dictionary suitable for direct use inside model computations.

The pipeline is: 1) Load payload (mapping or file-style behavior), 2) Canonicalize keys and fill defaults, 3) Enforce alias consistency, 4) Validate values and required fields.

Returns:

scaling_kwargs – Canonical and validated scaling configuration.

Return type:

Raises:

ValueError – If validation fails due to missing keys or invalid values.
KeyError – If canonicalization expects keys that are absent.
TypeError – If the payload contains unsupported types.

Notes

The returned dict is intended to be stable under Keras serialization and safe to store in model state.
This method always loads with copy=True to avoid mutating the stored payload.

Examples

>>> cfg = GeoPriorScalingConfig.from_any(
...     {"coords_normalized": True}
... )
>>> sk = cfg.resolve()
>>> sk["coords_normalized"]
True

See also

canonicalize_scaling_kwargs: Normalizes scaling keys and defaults.
validate_scaling_kwargs: Enforces schema and constraints.
enforce_scaling_alias_consistency: Prevents conflicting aliases.

get_config()[source]

Return a JSON-safe Keras configuration dictionary.

Keras uses this method to serialize the object. The returned dictionary must contain only JSON-serializable values.

This implementation uses _jsonify() to defensively convert nested structures such as NumPy scalars, tuples, and sets into plain Python types.

Returns:: config – JSON-safe configuration dictionary with the following keys: - payload: JSON-safe payload mapping, - source: provenance hint (may be None), - schema_version: schema version label.
Return type:: dict

Notes

source is stored for traceability and does not affect resolve().
When saved as part of a model config, this makes scaling reconstruction deterministic.

See also

GeoPriorScalingConfig.from_config: Recreate a config instance from this dictionary.

classmethod from_config(config)[source]

Recreate an instance from a Keras configuration dictionary.

This class method is used by Keras deserialization to rebuild the object from the dictionary returned by get_config().

Parameters:: config (dict) – Configuration dictionary produced by get_config().
Returns:: cfg – Reconstructed config instance.
Return type:: GeoPriorScalingConfig

Notes

This method does not call resolve(). Resolution is deferred to the consumer so that reconstruction remains explicit and testable.

See also

GeoPriorScalingConfig.get_config: Produces the configuration dictionary.

__init__(payload=<factory>, source=None, schema_version='1')

Parameters:

payload (dict)
source (str | None)
schema_version (str)

Return type:

None

geoprior.models.subsidence.scaling.override_scaling_kwargs(sk, cfg, *, finalize=None, dyn_names=None, gwl_dyn_index=None, base_dir=None, path_key='SCALING_KWARGS_JSON_PATH', strict=True, add_path=True, log_fn=None)[source]

Override scaling_kwargs from a JSON file or dict.

This helper applies an optional, precedence-based override to an existing scaling_kwargs mapping. The override source is read from cfg[path_key]. If the key is missing or empty, the input sk is returned (optionally finalized).

The override can be provided as:

a file path to a JSON object (mapping), or
a Python dict-like mapping embedded in cfg.

Overrides are applied via a deep-merge strategy:

for nested dict values, keys are merged recursively,
for non-dict values, the override replaces the base value.

Optionally, the merged result is passed through finalize to recompute derived or canonical fields (for example, coordinate ranges, unit flags, or other normalization metadata).

Parameters:

sk (Mapping[str, Any]) – Base scaling configuration (scaling_kwargs). This is typically computed by Stage-2 or loaded from Stage-1 output. The input is copied to a plain dict before modification.
cfg (Mapping[str, Any] or None) – Configuration mapping that may contain the override source under path_key. If None, no override is applied.
finalize (callable or None, optional) –
Function applied to the scaling dict to enforce canonical structure or to compute derived fields. If provided, it is applied before and after the override merge:
- pre-merge: normalize the base dict,
- post-merge: ensure the merged dict is consistent.
The callable must accept a dict and return a dict.
dyn_names (Sequence[str] or None, optional) – Expected dynamic feature names for safety validation. If provided and the override contains dynamic_feature_names, the two sequences are compared. A mismatch raises an error when strict=True.
gwl_dyn_index (int or None, optional) – Expected dynamic index for the groundwater-level feature. If provided and the override contains gwl_dyn_index, the values are compared. A mismatch raises an error when strict=True.
base_dir (str or None, optional) – Base directory used to resolve relative JSON paths. If None, the current working directory is used.
path_key (str, default "SCALING_KWARGS_JSON_PATH") – Name of the key in cfg that specifies the override. The value may be a dict-like mapping or a path to a JSON file.
strict (bool, default True) – Controls behavior on safety-check mismatches. When True, mismatches raise a ValueError. When False, mismatches can be logged via log_fn and the override still proceeds.
add_path (bool, default True) – If True, store the resolved override source in the output dict under scaling_kwargs_override_path. When the override is provided as a mapping (not a file), the value is set to "<dict>".
log_fn (callable or None, optional) – Optional logger function. If provided, it is called with informative messages such as successful override application and (when strict=False) mismatch warnings. Common choices are print or logger.info.

Returns:

out – Final scaling dict after optional override and optional finalization. The returned dict is independent from the input mapping object sk (a copy is always created).

Return type:

Raises:

FileNotFoundError – If cfg[path_key] is a path and the file does not exist.
ValueError – If a path is provided but the file does not contain valid JSON, or if a safety check fails while strict=True.
TypeError – If the loaded override is not a JSON object (dict-like).

Notes

Path resolution

When cfg[path_key] is a string path, it is resolved as:

Expand environment variables and ~.
If relative, join with base_dir (or CWD).

Safety checks

The checks are intentionally conservative. They prevent using an override file produced for a different dataset or feature layout. Recommended checks are:

dynamic_feature_names equality when known.
gwl_dyn_index equality when known.

You can extend validation by checking additional keys such as coord_epsg_used, coords_normalized, or unit flags.

Finalization In GeoPrior pipelines, finalize is typically a helper that enforces defaults and recomputes derived entries. Applying it both before and after the override helps reduce edge cases where the override only supplies partial information.

Figure assembly follows the plotting conventions described in Hunter [15].

Examples

Stage-2: override computed scaling with a file

In Stage-2, call this right after the auto-computed scaling is available, so the override takes precedence:

>>> sk = subsmodel_params["scaling_kwargs"]
>>> sk = override_scaling_kwargs(
...     sk,
...     cfg,
...     finalize=finalize_scaling_kwargs,
...     dyn_names=DYN_NAMES,
...     gwl_dyn_index=GWL_DYN_INDEX,
...     base_dir=os.path.dirname(__file__),
...     strict=True,
...     log_fn=print,
... )
>>> subsmodel_params["scaling_kwargs"] = sk

Stage-3: override Stage-1 scaling prior to enforcing bounds

In Stage-3, apply the override before injecting Stage-3 bounds:

>>> sk_model = dict(cfg.get("scaling_kwargs", {}) or {})
>>> sk_model = override_scaling_kwargs(
...     sk_model,
...     cfg,
...     dyn_names=sk_model.get("dynamic_feature_names"),
...     gwl_dyn_index=sk_model.get("gwl_dyn_index"),
...     base_dir=os.path.dirname(__file__),
... )
>>> sk_model["bounds"] = {
...     **(sk_model.get("bounds", {}) or {}),
...     **bounds_for_scaling,
... }

Inline dict override (no JSON file)

If the override is embedded in config, it is used directly:

>>> cfg = {
...     "SCALING_KWARGS_JSON_PATH": {
...         "coords_normalized": True,
...         "coord_ranges": {"t": 7.0, "x": 1000.0, "y": 900.0},
...     }
... }
>>> out = override_scaling_kwargs({}, cfg)

See also

finalize_scaling_kwargs: Canonicalize and complete scaling_kwargs entries.
compute_scaling_kwargs: Build a base scaling dict from data and pipeline settings.

Key scaling object#

class geoprior.models.subsidence.scaling.GeoPriorScalingConfig(payload=<factory>, source=None, schema_version='1')[source]

Bases: object

Scaling configuration utilities for GeoPrior PINN.

This module defines GeoPriorScalingConfig, a small Keras-serializable container used to store and reconstruct the physics scaling and slicing controls used by GeoPriorSubsNet.

The scaling configuration is critical because it governs how coordinates, time units, groundwater variables, and physics residuals are interpreted and non-dimensionalized. If this configuration is not faithfully serialized via Keras get_config(), a reloaded model may be reconstructed with a different effective physics behavior.

The main entry point is GeoPriorScalingConfig.from_any(), which accepts a dict-like mapping, a file path str, or an existing GeoPriorScalingConfig instance. The resolved configuration is produced by resolve(), which runs the same canonicalization and validation pipeline used during training.

Notes

The resolved scaling dictionary should be JSON-safe and stable under Keras serialization.
Use _jsonify() to defensively convert nested values (NumPy scalars, tuples, sets) into plain Python types.
The config container combines Keras serialization patterns with the standard-library dataclass model [16, 17].

See also

load_scaling_kwargs: Load scaling configuration from mapping or file.
canonicalize_scaling_kwargs: Normalize keys and fill defaults consistently.
enforce_scaling_alias_consistency: Ensure alias keys agree and do not conflict.
validate_scaling_kwargs: Validate schema and value ranges.

Parameters:

payload (dict)
source (str | None)
schema_version (str)

payload: dict

source: str | None = None

schema_version: str = '1'

classmethod from_any(obj, *, copy=True)[source]

Serializable container for GeoPrior scaling configuration.

This dataclass stores a “payload” dictionary that holds all scaling and physics-control parameters required to reproduce the model behavior after saving and reloading with Keras.

The container supports flexible construction from: - None (empty config), - a mapping (dict-like), - a file path str (loaded via load_scaling_kwargs), - an existing GeoPriorScalingConfig instance.

The canonical and validated configuration is produced by resolve(), which applies the GeoPrior scaling pipeline: loading, canonicalization, alias consistency checks, and validation.

Parameters:

payload (dict, optional) – Raw scaling configuration payload. This may be incomplete or contain aliases prior to canonicalization.
source (str or None, optional) – Optional provenance string, typically a file path used to load the payload. This is stored for traceability only.
schema_version (str, optional) – Version label for the payload schema. This can be used to implement migrations when the scaling format evolves.

Variables:

payload (dict) – The raw payload stored in this object.
source (str or None) – The provenance hint, if provided.
schema_version (str) – Schema version label.

Notes

The resolved scaling dictionary returned by resolve() is the one you should pass to the model internals.
get_config returns JSON-safe objects only. This avoids subtle reconstruction drift caused by non-serializable values.
This factory aligns with the Keras object-serialization pattern described in Keras Team [16].

Examples

Construct from a mapping:

>>> cfg = GeoPriorScalingConfig.from_any(
...     {"coords_normalized": True}
... )
>>> sk = cfg.resolve()
>>> isinstance(sk, dict)
True

Construct from a file path:

>>> cfg = GeoPriorScalingConfig.from_any(
...     "path/to/scaling_kwargs.json"
... )
>>> sk = cfg.resolve()

Use in a model constructor (pattern):

>>> cfg = GeoPriorScalingConfig.from_any(scaling_kwargs)
>>> scaling_kwargs_resolved = cfg.resolve()

See also

GeoPriorScalingConfig.from_any: Build config from dict, path, or config instance.
GeoPriorScalingConfig.resolve: Produce canonical and validated scaling dictionary.

load_scaling_kwargs, canonicalize_scaling_kwargs

resolve()[source]

Resolve the payload into a canonical, validated scaling dict.

This method runs the GeoPrior scaling pipeline and returns a dictionary suitable for direct use inside model computations.

The pipeline is: 1) Load payload (mapping or file-style behavior), 2) Canonicalize keys and fill defaults, 3) Enforce alias consistency, 4) Validate values and required fields.

Returns:

scaling_kwargs – Canonical and validated scaling configuration.

Return type:

Raises:

ValueError – If validation fails due to missing keys or invalid values.
KeyError – If canonicalization expects keys that are absent.
TypeError – If the payload contains unsupported types.

Notes

The returned dict is intended to be stable under Keras serialization and safe to store in model state.
This method always loads with copy=True to avoid mutating the stored payload.

Examples

>>> cfg = GeoPriorScalingConfig.from_any(
...     {"coords_normalized": True}
... )
>>> sk = cfg.resolve()
>>> sk["coords_normalized"]
True

See also

canonicalize_scaling_kwargs: Normalizes scaling keys and defaults.
validate_scaling_kwargs: Enforces schema and constraints.
enforce_scaling_alias_consistency: Prevents conflicting aliases.

get_config()[source]

Return a JSON-safe Keras configuration dictionary.

Keras uses this method to serialize the object. The returned dictionary must contain only JSON-serializable values.

This implementation uses _jsonify() to defensively convert nested structures such as NumPy scalars, tuples, and sets into plain Python types.

Returns:: config – JSON-safe configuration dictionary with the following keys: - payload: JSON-safe payload mapping, - source: provenance hint (may be None), - schema_version: schema version label.
Return type:: dict

Notes

source is stored for traceability and does not affect resolve().
When saved as part of a model config, this makes scaling reconstruction deterministic.

See also

GeoPriorScalingConfig.from_config: Recreate a config instance from this dictionary.

classmethod from_config(config)[source]

Recreate an instance from a Keras configuration dictionary.

This class method is used by Keras deserialization to rebuild the object from the dictionary returned by get_config().

Parameters:: config (dict) – Configuration dictionary produced by get_config().
Returns:: cfg – Reconstructed config instance.
Return type:: GeoPriorScalingConfig

Notes

This method does not call resolve(). Resolution is deferred to the consumer so that reconstruction remains explicit and testable.

See also

GeoPriorScalingConfig.get_config: Produces the configuration dictionary.

__init__(payload=<factory>, source=None, schema_version='1')

Parameters:

payload (dict)
source (str | None)
schema_version (str)

Return type:

None

Loss and step-packing helpers#

The losses module contains helpers used by the custom training and evaluation steps to pack structured outputs, including physics losses and diagnostics. In practice, this layer is where the public model outputs are reorganized into the richer objects consumed by the staged workflow.

GeoPrior loss assembly and logging helpers.

This module centralizes: - physics loss assembly (no double offset) - return packaging for train/test/eval

geoprior.models.subsidence.losses.should_log_physics(model)[source]

Decide whether to expose physics keys in logs.

If physics is off, logs are included only if scaling_kwargs[“log_physics_when_off”] is True.

Parameters:: model (Any)
Return type:: bool

geoprior.models.subsidence.losses.assemble_physics_loss(model, *, loss_cons, loss_gw, loss_prior, loss_smooth, loss_mv, loss_q_reg, loss_bounds)[source]

Assemble the full physics objective with an offset-aware multiplier.

This helper combines individual physics loss components computed by the GeoPrior PINN core into:

an unscaled physics loss (for logging and diagnostics),
a scaled physics loss (the quantity added to the data loss),
the global physics multiplier used for scaling,
a dictionary of per-term scaled contributions that is consistent with the scaled physics loss.

The function implements the GeoPrior weighting convention:

Each component loss is first multiplied by its corresponding per-term weight stored on the model instance (lambda_*).
A global physics multiplier phys_mult is computed by model._physics_loss_multiplier(), which depends on model.offset_mode and the scalar state model._lambda_offset.
The multiplier is applied to PDE-style terms by default, while certain calibration/regularization terms can opt out depending on model flags (see Notes).

Formally, define weighted terms:

(95)#\[\begin{split}T_{cons} = \lambda_{cons} L_{cons} \\ T_{gw} = \lambda_{gw} L_{gw} \\ T_{prior} = \lambda_{prior} L_{prior} \\ T_{smooth} = \lambda_{smooth} L_{smooth} \\ T_{bounds} = \lambda_{bounds} L_{bounds} \\ T_{mv} = \lambda_{mv} L_{mv} \\ T_{q} = \lambda_{q} L_{q}\end{split}\]

Let the PDE core sum be:

(96)#\[L_{core} = T_{cons} + T_{gw} + T_{prior} + T_{smooth} + T_{bounds}\]

and the unscaled physics loss be:

(97)#\[L_{phys,raw} = L_{core} + T_{mv} + T_{q}\]

The scaled physics loss is:

(98)#\[L_{phys,scaled} = phys\_mult \, L_{core} + s_{mv} \, T_{mv} + s_{q} \, T_{q}\]

where:

\(s_{mv} = phys\_mult\) if model._scale_mv_with_offset is True, else \(s_{mv} = 1\).
\(s_{q} = phys\_mult\) if model._scale_q_with_offset is True, else \(s_{q} = 1\).

Parameters:

model (Any) –
Model-like object providing the weighting attributes:
- lambda_cons, lambda_gw, lambda_prior, lambda_smooth, lambda_bounds, lambda_mv, lambda_q
- _physics_loss_multiplier() method
- optional flags _scale_mv_with_offset and _scale_q_with_offset
loss_cons (Tensor) – Consolidation loss \(L_{cons}\) (typically mean-square of a scaled consolidation residual).
loss_gw (Tensor) – Groundwater-flow PDE loss \(L_{gw}\) (typically mean-square of a scaled groundwater residual).
loss_prior (Tensor) – Timescale-consistency prior loss \(L_{prior}\) (often mean-square of a log-timescale residual).
loss_smooth (Tensor) – Smoothness prior loss \(L_{smooth}\) (regularizes spatial gradients of learned fields).
loss_mv (Tensor) – Storage identity / compressibility calibration loss \(L_{mv}\).
loss_q_reg (Tensor) – Forcing regularization loss \(L_{q}\) (typically mean-square of the SI forcing field \(Q\)).
loss_bounds (Tensor) – Soft-bounds penalty loss \(L_{bounds}\) derived from bound residuals (if enabled).

Returns:

physics_raw (Tensor) – Unscaled physics loss:

(99)#\[L_{phys,raw} = L_{core} + T_{mv} + T_{q}\]

Useful for diagnostics, independent of lambda_offset.
physics_scaled (Tensor) – Scaled physics loss, consistent with the global multiplier and the optional scaling rules for mv and q terms.
phys_mult (Tensor) – The global physics multiplier returned by model._physics_loss_multiplier().
terms_scaled (dict[str, Tensor]) – Per-term contributions consistent with physics_scaled. Keys are:
- 'cons', 'gw', 'prior', 'smooth', 'bounds', 'mv', 'q'.

Return type:

tuple[Tensor, Tensor, Tensor, dict[str, Tensor]]

Notes

Offset-aware scaling policy. The global multiplier phys_mult is intended as a single knob to warm up or damp all PDE-style physics terms together. By default:

PDE-style terms (cons, gw, prior, smooth, bounds) are always scaled by phys_mult.
The mv term is treated as a calibration loss and is not scaled by phys_mult unless model._scale_mv_with_offset is True.
The q regularization term is scaled by phys_mult only if model._scale_q_with_offset is True.

This separation avoids unintended suppression of calibration signals when physics warmup is used.

Logging and gradient debugging. Returning both physics_raw and physics_scaled helps debug training stability:

physics_raw shows whether residual magnitudes are decreasing.
physics_scaled shows the effective contribution to the total optimization objective.

The physics-informed weighting pattern follows Raissi et al. [18].

Examples

Assemble physics loss inside a training loop:

>>> physics_raw, physics_scaled, phys_mult, terms = (
...     assemble_physics_loss(
...         model,
...         loss_cons=loss_cons,
...         loss_gw=loss_gw,
...         loss_prior=loss_prior,
...         loss_smooth=loss_smooth,
...         loss_mv=loss_mv,
...         loss_q_reg=loss_q_reg,
...         loss_bounds=loss_bounds,
...     )
... )
>>> total_loss = data_loss + physics_scaled

Inspect per-term contributions:

>>> float(terms["prior"])
0.0123

See also

geoprior.models.subsidence.step_core.physics_core: Produces the component losses used as inputs here.
GeoPriorSubsNet.compile: Configures the lambda_* weights and the offset multiplier.

geoprior.models.subsidence.losses.zero_physics_bundle(model, *, dtype=tf.float32)[source]

Canonical zero physics bundle.

This keeps dashboards stable when requested.

Parameters:

model (Any)
dtype (Any)

Return type:

dict[str, Tensor]

geoprior.models.subsidence.losses.build_physics_bundle(model, *, physics_loss_raw, physics_loss_scaled, phys_mult, loss_cons, loss_gw, loss_prior, loss_smooth, loss_mv, loss_q_reg, q_rms, q_gate, subs_resid_gate, loss_bounds, eps_prior, eps_cons, eps_gw, eps_cons_raw=None, eps_gw_raw=None)[source]

Canonical physics bundle used by train/test/eval packers.

Parameters:

model (Any)
physics_loss_raw (Tensor)
physics_loss_scaled (Tensor)
phys_mult (Tensor)
loss_cons (Tensor)
loss_gw (Tensor)
loss_prior (Tensor)
loss_smooth (Tensor)
loss_mv (Tensor)
loss_q_reg (Tensor)
q_rms (Tensor)
q_gate (Tensor)
subs_resid_gate (Tensor)
loss_bounds (Tensor)
eps_prior (Tensor)
eps_cons (Tensor)
eps_gw (Tensor)
eps_cons_raw (Any | None)
eps_gw_raw (Any | None)

Return type:

dict[str, Tensor]

geoprior.models.subsidence.losses.update_epsilon_metrics(model, *, eps_prior, eps_cons, eps_gw)[source]

Update optional epsilon metrics if present.

Parameters:

model (Any)
eps_prior (Tensor)
eps_cons (Tensor)
eps_gw (Tensor)

Return type:

None

geoprior.models.subsidence.losses.epsilon_value_for_logs(model, which, fallback)[source]

Prefer tracked epsilon metric if it exists.

Parameters:

model (Any)
which (str)
fallback (Tensor)

Return type:

Tensor

geoprior.models.subsidence.losses.update_compiled_metrics(model, targets, y_pred)[source]

Update compiled Keras metrics for multi-output dict predictions.

This helper updates the metric container created by tf.keras.Model.compile() in a way that is robust across Keras 2 and Keras 3 behavior when the model uses named outputs (dict-style) and the training loop uses a custom train_step() / test_step().

The function:

Locates the “real” compiled metrics object for the model (if any) using an internal helper (_get_real_compile_metrics).
Determines the ordered list of output keys from the model (preferably model.output_names and then model._output_keys).
Aligns the shapes of ground truth tensors to match prediction tensors (via _as_BHO), so metrics always see consistent batch layout.
Attempts to update metrics using the most stable calling pattern for the installed Keras version:
- First try list-based update (update_state(y_true_list, y_pred_list)), which avoids dict key routing issues that can occur with certain Keras 2 configurations.
- If that fails, fall back to dict-based update (update_state(y_true_dict, y_pred_dict)).
- If that also fails, fall back to manually updating per-output metric objects by matching metric name prefixes.

This helper is primarily used to keep metric reporting consistent when custom training logic bypasses the default Keras fit loop internals.

Parameters:

model (Any) – A Keras model instance (or model-like object) that has been compiled with metrics and possibly multi-output losses.
targets (dict-like) – Ground truth outputs keyed by output name. Values can be tensors or tensor-like arrays.
y_pred (dict-like) – Model predictions keyed by output name. Values are tensors.

Returns:

Updates the compiled metrics state in-place.

Return type:

None

Notes

Why a custom updater is needed. Keras multi-output metric routing depends on how metrics were compiled (list-based vs dict-based) and how outputs are named and returned. In custom train_step() / test_step(), you often compute losses manually and must also call metric updates manually to preserve the behavior of model.fit.

Compatibility behavior. - In some Keras 2 environments, calling compiled.update_state with

dicts can fail or silently mis-route metrics when output names do not align with how the metric container was constructed. The list-first strategy is a defensive approach.

The final manual fallback updates metric objects directly by matching their name prefix (<output_name>_) and skipping loss-like metrics.

Shape normalization. The helper normalizes ground-truth shapes to match prediction shapes before updating metrics. This reduces common failures when targets are provided as (B,H) or (B,H,1) while predictions may be (B,H,Q,1) (quantiles) or similar.

Metric routing behavior follows Keras Team [19].

Examples

Inside a custom test_step:

>>> y_pred = model(inputs, training=False)
>>> update_compiled_metrics(model, targets, y_pred)

Inside a custom train_step:

>>> with tf.GradientTape() as tape:
...     y_pred = model(inputs, training=True)
...     loss = model.compiled_loss(...)
>>> update_compiled_metrics(model, targets, y_pred)

See also

tf.keras.Model.compiled_metrics: Standard entry point for metric containers in Keras.
GeoPriorSubsNet.train_step: Custom training loop that may use this helper to keep metrics consistent.

geoprior.models.subsidence.losses.safe_metric_result(metric, fallback=0.0)[source]

Safely obtain a metric result (Keras 3-safe).

In Keras 3, calling metric.result() may raise if the metric hasn’t been built/updated yet. In that case we return fallback.

Parameters:

metric (Any) – A Keras metric instance (or a scalar/tensor-like).
fallback (float, default 0.0) – Value returned if the metric is not ready or errors.

Returns:

Metric result as a float32 tensor (or fallback).

Return type:

Tensor

geoprior.models.subsidence.losses.pack_step_results(model, *, total_loss, data_loss, targets, y_pred, physics=None, manual_trackers=None)[source]

Canonical return dictionary for custom train_step / test_step.

This helper builds a stable logging payload for GeoPrior-style models that use a custom training loop. It combines:

supervised loss scalars (data and total),
compiled Keras metrics (if available),
optional manual trackers (e.g., add-on quantile trackers),
optional physics diagnostics (PINN losses and epsilons).

The function is intentionally defensive across Keras versions:

It explicitly updates and reads compiled metrics using update_compiled_metrics and the underlying compile-metrics container, rather than relying on model.metrics alone.
It reserves the key "loss" as the authoritative scalar returned to Keras, while also including explicit "total_loss" and "data_loss" entries for clarity.

Parameters:

model (Any) –
Model-like object that provides compiled metrics and configuration. Expected attributes and helpers include:
- metrics (optional list of metric objects)
- output_names or _output_keys (output ordering)
- scaling_kwargs (optional dict)
- functions used by this module such as should_log_physics, zero_physics_bundle, update_compiled_metrics, safe_metric_result, update_epsilon_metrics, and epsilon_value_for_logs.
total_loss (Tensor) – The scalar loss used for optimization in the current step. This is returned as results["loss"] and results["total_loss"].
data_loss (Tensor) – The supervised loss computed from the compiled loss function (i.e., the data term). Returned as results["data_loss"].
targets (Any) – Ground-truth targets for the supervised outputs. Typically a dict keyed by output names (e.g., {"subs_pred": ..., "gwl_pred": ...}) but may be any structure supported by update_compiled_metrics.
y_pred (Any) – Predicted outputs corresponding to targets. Typically a dict keyed by output names.
physics (dict[str, Tensor] or None, optional) – Physics bundle produced by physics_core (or an equivalent). If None and physics logging is enabled, a zero bundle is used.
manual_trackers (dict or None, optional) – Optional additional trackers to log. Values may be metric objects with result() or raw scalars/tensors. This is typically used for add-on metrics that are not part of Keras compiled metrics.

Returns:

results – A dictionary suitable for returning from train_step or test_step. At minimum it contains:

loss: total loss used by Keras progress reporting.
total_loss: same as loss (explicit alias).
data_loss: supervised/data loss term.

If compiled metrics are available, additional keys are included (e.g., subs_pred_mae, quantile coverage, etc.). If physics logging is enabled, physics diagnostics are appended (see Notes).

Return type:

dict[str, Tensor]

Notes

Metric collection strategy. Compiled metrics are updated via update_compiled_metrics and then read from the underlying compile-metrics object. This avoids common routing failures when using dict outputs in custom training loops.

Reserved and excluded keys. Certain names are reserved to prevent collisions with Keras internals and to ensure that the loss scalar remains authoritative. Some epsilon fields may also be excluded from the compiled-metric collection to avoid duplicate/conflicting reporting.

Physics logging. If physics logging is enabled (should_log_physics(model) returns True), this helper adds a consistent set of physics metrics, typically:

physics losses (raw and scaled),
per-term losses (consolidation, gw flow, priors, bounds),
epsilon metrics (scaled and raw variants).

If physics is disabled for the model and logging is enabled, a zero bundle is inserted to keep log schemas stable.

Q and residual gates. When scaling_kwargs requests Q diagnostics (log_q_diagnostics=True), additional fields such as Q RMS and gate values may be included for debugging training schedules.

The custom-loop packing pattern follows Keras Team [20].

Examples

Inside a custom training step:

>>> results = pack_step_results(
...     model,
...     total_loss=total_loss,
...     data_loss=data_loss,
...     targets=targets,
...     y_pred=y_pred,
...     manual_trackers=(model.add_on.as_dict if model.add_on else None),
...     physics=physics_bundle,
... )
>>> return results

Inside a custom test step:

>>> return pack_step_results(
...     model,
...     total_loss=total_loss,
...     data_loss=data_loss,
...     targets=targets,
...     y_pred=y_pred,
...     physics=physics_bundle,
... )

See also

update_compiled_metrics: Compatibility helper to update metrics for multi-output dicts.
assemble_physics_loss: Builds the scaled physics objective used in total_loss.
physics_core: Produces the physics bundle consumed by this packer.

geoprior.models.subsidence.losses.pack_eval_physics(model, *, physics)[source]

Canonical physics bundle output for batch-level physics evaluation.

This helper normalizes the output of physics diagnostics so that callers can rely on a stable schema regardless of whether physics is enabled for the model.

Behavior:

If a physics bundle is provided, it is returned unchanged.
If physics is off and logging is enabled, a zero-valued physics bundle is returned (to keep downstream logging stable).
If physics is off and logging is disabled, an empty dict is returned.

Parameters:

model (Any) – Model-like object that controls whether physics logging is enabled. This function relies on should_log_physics(model) and zero_physics_bundle(model) which are expected to be available in the surrounding module.
physics (dict[str, Tensor] or None) – Physics bundle produced by physics_core or a compatible routine. If None, behavior depends on whether physics logging is enabled.

Returns:

out – Canonical physics dictionary.

If physics is enabled (or logging when off), keys typically include (implementation dependent):

physics_loss_raw
physics_loss_scaled
physics_mult
per-term losses and epsilon diagnostics

If physics is off and logging is disabled, returns {}.

Return type:

dict[str, Tensor]

Notes

Returning a zero bundle when physics is off is useful for dashboards and automated training loops where missing keys complicate aggregation.

Examples

Batch-level evaluation:

>>> packed = pack_eval_physics(model, physics=physics_bundle)

Physics-off scenario:

>>> packed = pack_eval_physics(model, physics=None)
>>> packed  # either {} or a zero bundle depending on settings

See also

GeoPriorSubsNet.evaluate_physics: Aggregates these batch outputs across datasets.
physics_core: Produces the physics bundle consumed by this helper.

Important helpers#

geoprior.models.subsidence.losses.pack_step_results(model, *, total_loss, data_loss, targets, y_pred, physics=None, manual_trackers=None)[source]

Canonical return dictionary for custom train_step / test_step.

This helper builds a stable logging payload for GeoPrior-style models that use a custom training loop. It combines:

supervised loss scalars (data and total),
compiled Keras metrics (if available),
optional manual trackers (e.g., add-on quantile trackers),
optional physics diagnostics (PINN losses and epsilons).

The function is intentionally defensive across Keras versions:

It explicitly updates and reads compiled metrics using update_compiled_metrics and the underlying compile-metrics container, rather than relying on model.metrics alone.
It reserves the key "loss" as the authoritative scalar returned to Keras, while also including explicit "total_loss" and "data_loss" entries for clarity.

Parameters:

model (Any) –
Model-like object that provides compiled metrics and configuration. Expected attributes and helpers include:
- metrics (optional list of metric objects)
- output_names or _output_keys (output ordering)
- scaling_kwargs (optional dict)
- functions used by this module such as should_log_physics, zero_physics_bundle, update_compiled_metrics, safe_metric_result, update_epsilon_metrics, and epsilon_value_for_logs.
total_loss (Tensor) – The scalar loss used for optimization in the current step. This is returned as results["loss"] and results["total_loss"].
data_loss (Tensor) – The supervised loss computed from the compiled loss function (i.e., the data term). Returned as results["data_loss"].
targets (Any) – Ground-truth targets for the supervised outputs. Typically a dict keyed by output names (e.g., {"subs_pred": ..., "gwl_pred": ...}) but may be any structure supported by update_compiled_metrics.
y_pred (Any) – Predicted outputs corresponding to targets. Typically a dict keyed by output names.
physics (dict[str, Tensor] or None, optional) – Physics bundle produced by physics_core (or an equivalent). If None and physics logging is enabled, a zero bundle is used.
manual_trackers (dict or None, optional) – Optional additional trackers to log. Values may be metric objects with result() or raw scalars/tensors. This is typically used for add-on metrics that are not part of Keras compiled metrics.

Returns:

results – A dictionary suitable for returning from train_step or test_step. At minimum it contains:

loss: total loss used by Keras progress reporting.
total_loss: same as loss (explicit alias).
data_loss: supervised/data loss term.

If compiled metrics are available, additional keys are included (e.g., subs_pred_mae, quantile coverage, etc.). If physics logging is enabled, physics diagnostics are appended (see Notes).

Return type:

dict[str, Tensor]

Notes

Metric collection strategy. Compiled metrics are updated via update_compiled_metrics and then read from the underlying compile-metrics object. This avoids common routing failures when using dict outputs in custom training loops.

Reserved and excluded keys. Certain names are reserved to prevent collisions with Keras internals and to ensure that the loss scalar remains authoritative. Some epsilon fields may also be excluded from the compiled-metric collection to avoid duplicate/conflicting reporting.

Physics logging. If physics logging is enabled (should_log_physics(model) returns True), this helper adds a consistent set of physics metrics, typically:

physics losses (raw and scaled),
per-term losses (consolidation, gw flow, priors, bounds),
epsilon metrics (scaled and raw variants).

If physics is disabled for the model and logging is enabled, a zero bundle is inserted to keep log schemas stable.

Q and residual gates. When scaling_kwargs requests Q diagnostics (log_q_diagnostics=True), additional fields such as Q RMS and gate values may be included for debugging training schedules.

The custom-loop packing pattern follows Keras Team [20].

Examples

Inside a custom training step:

>>> results = pack_step_results(
...     model,
...     total_loss=total_loss,
...     data_loss=data_loss,
...     targets=targets,
...     y_pred=y_pred,
...     manual_trackers=(model.add_on.as_dict if model.add_on else None),
...     physics=physics_bundle,
... )
>>> return results

Inside a custom test step:

>>> return pack_step_results(
...     model,
...     total_loss=total_loss,
...     data_loss=data_loss,
...     targets=targets,
...     y_pred=y_pred,
...     physics=physics_bundle,
... )

See also

update_compiled_metrics: Compatibility helper to update metrics for multi-output dicts.
assemble_physics_loss: Builds the scaled physics objective used in total_loss.
physics_core: Produces the physics bundle consumed by this packer.

Identifiability controls#

The identifiability layer exposes regime setup, compile-time weight resolution, head locks, and audit helpers. This is an important part of the scientific design of GeoPrior-v3, because it controls how strongly different fields are tied to their priors and which degrees of freedom remain open during training.

Identifiability scenarios for GeoPrior-style models.

Goal: - break non-identifiability ridges by construction.

Option A: - learn tau only - derive K from tau via closure - freeze (or fix) Ss and Hd

geoprior.models.subsidence.identifiability.init_identifiability(regime, scaling_kwargs)[source]

Apply identifiability profile to scaling kwargs.

does NOT override user-provided keys
ensures sk[“bounds_loss”] exists (dict form)

Parameters:

regime (str | None)
scaling_kwargs (Mapping[str, Any])

Return type:

tuple[str | None, dict[str, Any] | None, dict]

geoprior.models.subsidence.identifiability.apply_ident_locks(model, profile=None)[source]

Parameters:

model (Any)
profile (Mapping[str, Any] | None)

Return type:

None

geoprior.models.subsidence.identifiability.resolve_compile_weights(profile, *, lambda_cons, lambda_gw, lambda_prior, lambda_smooth, lambda_mv, lambda_bounds, lambda_q)[source]

Parameters:

profile (Mapping[str, Any] | None)
lambda_cons (float | None)
lambda_gw (float | None)
lambda_prior (float | None)
lambda_smooth (float | None)
lambda_mv (float | None)
lambda_bounds (float | None)
lambda_q (float | None)

Return type:

dict[str, float]

geoprior.models.subsidence.identifiability.get_ident_profile(regime)[source]

Parameters:: regime (str | None)
Return type:: tuple[str | None, dict[str, Any] | None]

geoprior.models.subsidence.identifiability.ident_audit_dict(model, *, extra_sk_keys=None)[source]

Small, JSON-safe audit of identifiability configuration.

Intended for experiment logs / manifests / eval JSON.

Parameters:

model (Any)
extra_sk_keys (Iterable[str] | None)

Return type:

Important helpers#

geoprior.models.subsidence.identifiability.init_identifiability(regime, scaling_kwargs)[source]

Apply identifiability profile to scaling kwargs.

does NOT override user-provided keys
ensures sk[“bounds_loss”] exists (dict form)

Parameters:

regime (str | None)
scaling_kwargs (Mapping[str, Any])

Return type:

tuple[str | None, dict[str, Any] | None, dict]

geoprior.models.subsidence.identifiability.apply_ident_locks(model, profile=None)[source]

Parameters:

model (Any)
profile (Mapping[str, Any] | None)

Return type:

None

geoprior.models.subsidence.identifiability.resolve_compile_weights(profile, *, lambda_cons, lambda_gw, lambda_prior, lambda_smooth, lambda_mv, lambda_bounds, lambda_q)[source]

Parameters:

profile (Mapping[str, Any] | None)
lambda_cons (float | None)
lambda_gw (float | None)
lambda_prior (float | None)
lambda_smooth (float | None)
lambda_mv (float | None)
lambda_bounds (float | None)
lambda_q (float | None)

Return type:

dict[str, float]

Payloads and exported physics artifacts#

The payload layer provides helpers for gathering, saving, loading, and subsampling physics payloads used later in diagnostics, inference, and figure generation. These helpers matter for reproducibility because they make it possible to inspect the internal physical state of a run after training.

Physics diagnostics payloads.

This module centralizes data collection from a trained model for physics sanity plots (e.g., Fig.4) and provides robust persistence to disk with simple provenance metadata.

geoprior.models.subsidence.payloads.default_meta_from_model(model)[source]

Build lightweight, JSON-serializable provenance from a model.

Notes

time_units describes the time coordinate units in the dataset (for example, "year"), meaning what t represents before conversion.
Physics diagnostics such as tau, tau_prior, K, and cons_res_vals are exported in SI time units after the model’s internal conversions. In practice, K is in m/s, tau is in s, and cons_res_vals is in m/s.

Return type:: dict

geoprior.models.subsidence.payloads.identifiability_diagnostics_from_payload(payload, tau_true, K_true, Ss_true, Hd_true, K_prior, Ss_prior, Hd_prior, quantiles=(0.5, 0.75, 0.9, 0.95), eps=1e-12)[source]

Compute synthetic identifiability diagnostics from a physics payload.

This implements the three diagnostics described in Supplementary Methods 3:

Relative error in the effective relaxation time tau.

Discrepancy between the composite timescale closure H_d^2 S_s / (kappa K) (stored as tau_prior) and the true effective timescale tau_eff,true, via a log-timescale residual.

Marginal log-offsets of K, S_s and H_d relative to their true effective values and lithology-based priors.

Parameters:

payloaddict: Physics payload returned by gather_physics_payload() or GeoPriorSubsNet.export_physics_payload(). Must contain 1D arrays with keys: “tau”, “tau_prior”, “K”, “Ss”, “Hd”.
tau_truefloat: True effective relaxation time :math:` au_{mathrm{eff,true}}` from the 1D consolidation column.
K_true, Ss_true, Hd_truefloat: True effective closures \(K_{\mathrm{eff}}\), \(S_{s,\mathrm{eff}}\), and \(H_{d,\mathrm{eff}}\) at the column scale.
K_prior, Ss_prior, Hd_priorfloat: Lithology-based priors used to construct the GeoPrior head for this synthetic column.
quantilestuple of float, default (0.5, 0.75, 0.9, 0.95): Quantile levels used for summary statistics of the distributions.
epsfloat, default 1e-12: Lower bound used to clip strictly positive quantities before taking logarithms.

Returns:

dict

A dictionary with three blocks:

"tau_rel_error": statistics of the relative error :math:`

rac{| au - au_{true}|}{ au_{true}}`.

"closure_log_resid": statistics of the log-timescale residual log(tau_prior) - log(tau_true).
"offsets": nested dict with "vs_true" and "vs_prior", each containing summary stats for the log-offsets delta_K, delta_Ss, and delta_Hd.

Parameters:

payload (dict[str, ndarray])
tau_true (float)
K_true (float)
Ss_true (float)
Hd_true (float)
K_prior (float)
Ss_prior (float)
Hd_prior (float)
quantiles (tuple[float, ...])
eps (float)

Return type:

geoprior.models.subsidence.payloads.summarise_effective_params(payload)[source]

Collapse 1D arrays to scalar effective parameters.

Intended for 1D synthetic-column experiments where model outputs are spatially constant and we only need a single representative value per run.

Parameters:: payload (dict[str, ndarray])
Return type:: dict[str, float]

geoprior.models.subsidence.payloads.compute_identifiability_summary(eff_params, true_params, prior_params, kappa_b=1.0, eps=1e-12)[source]

Compute identifiability diagnostics for Supp. Methods 3.

See Supplementary Methods 3 for definitions of the quantities returned.

Parameters:

eff_params (dict[str, float])
true_params (dict[str, float])
prior_params (dict[str, float])
kappa_b (float)
eps (float)

Return type:

dict[str, float]

geoprior.models.subsidence.payloads.gather_physics_payload(model, dataset, max_batches=None, float_dtype=<class 'numpy.float32'>, log_fn=None, eps=1e-12, **tqdm_kws)[source]

Collect a flat physics payload from a batched dataset for diagnostics.

This function iterates over a tf.data.Dataset (or any iterable) and calls model.evaluate_physics(inputs, return_maps=True) on each batch. The returned per-batch tensors are flattened and concatenated into 1D arrays suitable for scatter plots, histograms, and summary stats.

Important

No unit conversion is performed here. The payload is exported in whatever units evaluate_physics(…) returns. Unit consistency is therefore a responsibility of the model’s physics implementation (and its scaling_kwargs), not this I/O layer.

Parameters:

dataset (Iterable)
max_batches (int | None)
eps (float)

Return type:

dict[str, ndarray]

geoprior.models.subsidence.payloads.save_physics_payload(payload, meta, path=None, format='npz', overwrite=False, log_fn=None)[source]

Save payload + sidecar metadata to disk.

Parameters:

payload (dict) – Output of gather_physics_payload.
meta (dict) – Provenance dictionary. Will be JSON-serialized.
path (str or Nonr) – File path. If extension missing, inferred from format. If not provided, then get the current directory instead.
format ({"npz","csv","parquet"}) – Storage format. “npz” is compact and dependency-free.
overwrite (bool) – If False, raise if the file already exists.

Returns:

The resolved data file path that was written.

Return type:

str

geoprior.models.subsidence.payloads.load_physics_payload(path)[source]

Load a previously saved physics payload and its metadata.

Parameters:: path (str) – Data file path. Supports .npz, .csv, .parquet.
Returns:: (payload, meta) – Payload dict with arrays and metadata dict (if found).
Return type:: (dict, dict)

Important helpers#

geoprior.models.subsidence.payloads.gather_physics_payload(model, dataset, max_batches=None, float_dtype=<class 'numpy.float32'>, log_fn=None, eps=1e-12, **tqdm_kws)[source]

Collect a flat physics payload from a batched dataset for diagnostics.

This function iterates over a tf.data.Dataset (or any iterable) and calls model.evaluate_physics(inputs, return_maps=True) on each batch. The returned per-batch tensors are flattened and concatenated into 1D arrays suitable for scatter plots, histograms, and summary stats.

Important

No unit conversion is performed here. The payload is exported in whatever units evaluate_physics(…) returns. Unit consistency is therefore a responsibility of the model’s physics implementation (and its scaling_kwargs), not this I/O layer.

Parameters:

dataset (Iterable)
max_batches (int | None)
eps (float)

Return type:

dict[str, ndarray]

geoprior.models.subsidence.payloads.save_physics_payload(payload, meta, path=None, format='npz', overwrite=False, log_fn=None)[source]

Save payload + sidecar metadata to disk.

Parameters:

payload (dict) – Output of gather_physics_payload.
meta (dict) – Provenance dictionary. Will be JSON-serialized.
path (str or Nonr) – File path. If extension missing, inferred from format. If not provided, then get the current directory instead.
format ({"npz","csv","parquet"}) – Storage format. “npz” is compact and dependency-free.
overwrite (bool) – If False, raise if the file already exists.

Returns:

The resolved data file path that was written.

Return type:

str

geoprior.models.subsidence.payloads.load_physics_payload(path)[source]

Load a previously saved physics payload and its metadata.

Parameters:: path (str) – Data file path. Supports .npz, .csv, .parquet.
Returns:: (payload, meta) – Payload dict with arrays and metadata dict (if found).
Return type:: (dict, dict)

Stability helpers#

The stability layer contains helper utilities used to keep training and evaluation robust in the presence of stiff physics branches, unstable residual terms, or bad gradients.

Numerical stability helpers for subsidence physics workflows.

geoprior.models.subsidence.stability.clamp_physics_logits(K_logits, Ss_logits, dlogtau_logits, Q_logits=None, clip_min=-15.0, clip_max=15.0)[source]

(Fix A) Clamps raw logits to prevent exponential explosion/underflow in the physics layer.

Range [-15, 15] corresponds to exp(-15) ~ 3e-7 and exp(15) ~ 3e6.

geoprior.models.subsidence.stability.sanitize_scales(scales, min_scale=1e-06, max_scale=1000000.0)[source]: (Fix B) Replaces NaN/Inf values in dynamic scaling factors with 1.0 and clamps extreme values.

geoprior.models.subsidence.stability.compute_physics_warmup_gate(step_tensor, warmup_steps=500, ramp_steps=500)[source]

(Fix C) Returns a scalar 0.0 -> 1.0 multiplier based on global step.

Logic:

step < warmup: 0.0
warmup < step < warmup+ramp: linear ramp 0->1
step > warmup+ramp: 1.0

geoprior.models.subsidence.stability.filter_nan_gradients(grads)[source]

Important helpers#

geoprior.models.subsidence.stability.filter_nan_gradients(grads)[source]

Utility helpers#

The utils module contains conversion and convenience helpers used across the subsidence stack, including SI mapping, groundwater/head conversion, initialization helpers, and policy gating.

GeoPrior subsidence model utilities.

geoprior.models.subsidence.utils.enforce_scaling_alias_consistency(scaling_kwargs, *, where='validate')[source]

Enforce that canonical keys and aliases agree.

If both canonical and an alias exist and their values differ, apply the scaling error policy.

Parameters:

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

Return type:

None

geoprior.models.subsidence.utils.canonicalize_scaling_kwargs(scaling_kwargs, *, copy=True)[source]

Return a canonicalized scaling dict.

If a canonical key is missing, but one of its aliases exists, copy alias -> canonical.
Keeps existing canonical values unchanged.

Parameters:

scaling_kwargs (dict[str, Any] | None)
copy (bool)

Return type:

geoprior.models.subsidence.utils.load_scaling_kwargs(scaling_kwargs, *, copy=True)[source]

Load scaling kwargs from a dict-like object or JSON.

Parameters:

scaling_kwargs (Any | None)
copy (bool)

Return type:

geoprior.models.subsidence.utils.get_sk(scaling_kwargs, key, *aliases, default=None, required=False, cast=None)[source]

Fetch a key from scaling_kwargs with aliases + default.

Tries: key -> built-in aliases -> explicit aliases
Treats None and blank strings as “missing” and keeps searching.

Parameters:

key (str)
aliases (str)
required (bool)

geoprior.models.subsidence.utils.validate_scaling_kwargs(scaling_kwargs)[source]

Basic scaling sanity checks.

This includes policy-controlled heuristic checks for common “silent fallback” cases.

Parameters:: scaling_kwargs (dict[str, Any] | None)
Return type:: None

geoprior.models.subsidence.utils.affine_from_cfg(scaling_kwargs, *, scale_key, bias_key, meta_keys=(), unit_key=None)[source]

Return (a,b) for y_si = y_model*a + b.

Parameters:

scaling_kwargs (dict[str, Any] | None)
scale_key (str)
bias_key (str)
meta_keys (tuple[str, ...])
unit_key (str | None)

Return type:

tuple[Tensor, Tensor]

geoprior.models.subsidence.utils.to_si_thickness(H_model, scaling_kwargs)[source]

Convert thickness to SI.

Parameters:

H_model (Tensor)
scaling_kwargs (dict[str, Any] | None)

Return type:

Tensor

geoprior.models.subsidence.utils.to_si_head(h_model, scaling_kwargs)[source]

Convert head/depth to SI meters.

Parameters:

h_model (Tensor)
scaling_kwargs (dict[str, Any] | None)

Return type:

Tensor

geoprior.models.subsidence.utils.to_si_subsidence(s_model, scaling_kwargs)[source]

Convert subsidence to SI meters.

Parameters:

s_model (Tensor)
scaling_kwargs (dict[str, Any] | None)

Return type:

Tensor

geoprior.models.subsidence.utils.from_si_subsidence(s_si, scaling_kwargs)[source]

Inverse of to_si_subsidence: s_model = (s_si - b) / a.

Parameters:

s_si (Tensor)
scaling_kwargs (dict[str, Any] | None)

Return type:

Tensor

geoprior.models.subsidence.utils.deg_to_m(axis, scaling_kwargs)[source]

Meters per degree factor for lon/lat coords.

If coords_in_degrees=True and deg_to_m_lon/lat are missing, we try to compute them from lat0_deg (recommended).

Parameters:

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

Return type:

Tensor

geoprior.models.subsidence.utils.coord_ranges(scaling_kwargs)[source]

Return (tR,xR,yR) if coords_normalized.

Parameters:: scaling_kwargs (dict[str, Any] | None)
Return type:: tuple[float | None, float | None, float | None]

geoprior.models.subsidence.utils.resolve_gwl_dyn_index(scaling_kwargs)[source]

Resolve GWL channel index for dynamic_features.

Parameters:: scaling_kwargs (dict[str, Any] | None)
Return type:: int

geoprior.models.subsidence.utils.get_gwl_dyn_index_cached(model)[source]

Cache gwl_dyn_index on model after first resolve.

Return type:: int

geoprior.models.subsidence.utils.resolve_subs_dyn_index(scaling_kwargs)[source]

Resolve subsidence channel index for dynamic_features.

This is optional: v3.2 can use historical subsidence as a dynamic driver to provide a physics-friendly initial condition for the mean settlement path.

geoprior.models.subsidence.utils.get_subs_dyn_index_cached(model)[source]

Cache subs_dyn_index on model after first resolve.

Return type:: int

geoprior.models.subsidence.utils.slice_dynamic_channel(Xh, idx)[source]

Slice (B,T,F) -> (B,T,1) at idx.

Parameters:

Xh (Tensor)
idx (int)

Return type:

Tensor

geoprior.models.subsidence.utils.assert_dynamic_names_match_tensor(Xh, scaling_kwargs)[source]

Check dynamic_feature_names length matches Xh.

Parameters:

Xh (Tensor)
scaling_kwargs (dict[str, Any] | None)

Return type:

None

geoprior.models.subsidence.utils.gwl_to_head_m(v_m, scaling_kwargs, *, inputs=None)[source]

Convert depth-bgs to head if possible.

Parameters:

v_m (Tensor)
scaling_kwargs (dict[str, Any] | None)
inputs (dict[str, Tensor] | None)

Return type:

Tensor

geoprior.models.subsidence.utils.get_h_hist_si(model, inputs, *, want_head=True)[source]

Return head (or depth) history in SI meters.

Parameters:

model (object) – The model instance (provides scaling_kwargs and cached indices).
inputs (dict) – Batch inputs; expects dynamic_features unless an explicit head history key is provided.
want_head (bool, default True) – If True, convert depth-bgs to hydraulic head when possible.

Returns:

(B,T,1) tensor in SI meters.

Return type:

Tensor

geoprior.models.subsidence.utils.get_s_init_si(model, inputs, like)[source]

Return initial settlement (cumulative subsidence) in SI meters.

Priority: 1) explicit keys in inputs (s_init_si/subs_hist_last_si/…) 2) last historical value from dynamic_features if subs_dyn_index exists 3) zeros (broadcast)

Parameters:

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

Return type:

Tensor

geoprior.models.subsidence.utils.get_h_ref_si(model, inputs, like)[source]

Return h_ref in SI meters, broadcast to like.

Parameters:

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

Return type:

Tensor

geoprior.models.subsidence.utils.infer_dt_units_from_t(t_BH1, scaling_kwargs, *, eps=1e-12)[source]

Infer per-step dt in time_units from time tensor t(B,H,1).

Parameters:

t_BH1 (Tensor)
scaling_kwargs (dict[str, Any] | None)
eps (float)

Return type:

Tensor

geoprior.models.subsidence.utils.policy_gate(step, policy, *, warmup_steps=0, ramp_steps=0, dtype=tf.float32)[source]

Return a scalar gate in [0,1] based on a policy + step.

Parameters:

step (Tensor) – Global step counter (typically optimizer.iterations).
policy ({"always_on","always_off","warmup_off"}) – Gating behavior. always_on returns 1, always_off returns 0, and warmup_off returns 0 for step < warmup_steps before ramping to 1 over ramp_steps when ramp_steps > 0 or switching immediately at warmup_steps otherwise.
warmup_steps (int, default 0) – Number of steps to keep the gate at 0 (only for warmup_off).
ramp_steps (int, default 0) – Number of steps for a linear ramp from 0->1 after warmup. If 0, the gate is a hard step.
dtype (dtype, default tf_float32) – Output dtype.

Return type:

Tensor

geoprior.models.subsidence.utils.finalize_scaling_kwargs(sk)[source]

Add derived SI conversion constants to scaling_kwargs.

Adds (when possible): - seconds_per_time_unit: float - coord_ranges_si: dict with keys t (seconds), x/y (meters) - coord_inv_ranges_si: inverse of the above (safe floor).

Notes

This helper is designed to be called once when assembling scaling_kwargs (e.g., in your stage2 script) so the model can reuse those constants without recomputing unit conversions in the hot training loop.

Parameters:: sk (dict[str, Any])
Return type:: dict[str, Any]

geoprior.models.subsidence.utils.coord_ranges_si(sk)[source]

Return coordinate spans in SI (t in seconds; x/y in meters).

If coord_ranges_si is present in sk, it is used directly. Otherwise, this is computed from coord_ranges and time_units (and degree-to-meter factors when applicable).

Parameters:: sk (dict[str, Any])
Return type:: tuple[float | None, float | None, float | None]

Selected public helpers#

geoprior.models.subsidence.utils.to_si_head(h_model, scaling_kwargs)[source]

Convert head/depth to SI meters.

Parameters:

h_model (Tensor)
scaling_kwargs (dict[str, Any] | None)

Return type:

Tensor

geoprior.models.subsidence.utils.to_si_thickness(H_model, scaling_kwargs)[source]

Convert thickness to SI.

Parameters:

H_model (Tensor)
scaling_kwargs (dict[str, Any] | None)

Return type:

Tensor

geoprior.models.subsidence.utils.from_si_subsidence(s_si, scaling_kwargs)[source]

Inverse of to_si_subsidence: s_model = (s_si - b) / a.

Parameters:

s_si (Tensor)
scaling_kwargs (dict[str, Any] | None)

Return type:

Tensor

geoprior.models.subsidence.utils.gwl_to_head_m(v_m, scaling_kwargs, *, inputs=None)[source]

Convert depth-bgs to head if possible.

Parameters:

v_m (Tensor)
scaling_kwargs (dict[str, Any] | None)
inputs (dict[str, Tensor] | None)

Return type:

Tensor

geoprior.models.subsidence.utils.get_h_ref_si(model, inputs, like)[source]

Return h_ref in SI meters, broadcast to like.

Parameters:

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

Return type:

Tensor

geoprior.models.subsidence.utils.get_s_init_si(model, inputs, like)[source]

Return initial settlement (cumulative subsidence) in SI meters.

Priority: 1) explicit keys in inputs (s_init_si/subs_hist_last_si/…) 2) last historical value from dynamic_features if subs_dyn_index exists 3) zeros (broadcast)

Parameters:

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

Return type:

Tensor

Shared physics core#

The step_core module contains the shared physics-core execution path that assembles the structured physics bundle used in training and evaluation. It is especially useful when you want to trace how model outputs, scaling metadata, physics residuals, and diagnostics are joined together in one place.

Core step computations for subsidence physics evaluation.

geoprior.models.subsidence.step_core.physics_core(model, inputs, training, return_maps=False, *, for_train=False)[source]

Compute GeoPrior physics residuals and losses for a batch.

This function implements the shared physics pathway used by both training and evaluation for GeoPrior-style PINN models. It is designed to keep the physics logic consistent across:

train_step() (when physics losses are added to the total loss)
evaluation routines (when physics diagnostics are reported)

At a high level, the function performs:

Input preparation and SI conversions (thickness, head, coords).
Forward pass through the model to obtain data predictions and physics logits.
Mapping of physics logits to bounded physical fields (\(K\), \(S_s\), \(tau\)) and the closure prior \(tau_{phys}\).
Automatic differentiation to obtain PDE derivatives with respect to the model coords.
Chain-rule scaling to SI-consistent derivatives.
Construction of residual maps for: * consolidation relaxation residual, * groundwater flow residual, * time-scale prior residual, * smoothness prior residual, * bounds residual.
Optional nondimensionalization / residual scaling.
Assembly of physics losses, gating schedules, and diagnostic epsilon metrics.

The returned dictionary contains predictions, auxiliary forward outputs, packed physics values (for logging), and optionally the full residual maps and fields.

Parameters:

model (object) –
Model instance providing GeoPrior-style methods and attributes.

The function expects the model to expose (at minimum):
- scaling_kwargsdict
  Resolved scaling and convention payload.
- time_unitsstr or None
  Dataset time unit (for per-second conversions).
- forecast_horizonint
  Horizon length used to tile coords when needed.
- _forward_all(inputs, training=...)callable
  Forward pass returning (y_pred, aux).
- split_data_predictions(x)callable
  Split concatenated data head into subsidence and GWL.
- split_physics_predictions(x)callable
  Split concatenated physics head into (K_logits, Ss_logits, dlogtau_logits, Q_logits).
- pde_modes_activeiterable of str
  Active PDE modes (e.g., {‘consolidation’, ‘gw_flow’}).
- Optional gates: _q_gate(), _subs_resid_gate().
- Optional physics switch: _physics_off().
The function is tolerant to partial capabilities and will short-circuit when physics is disabled, but missing mandatory signals (e.g., thickness) raise errors.
inputs (dict[str, Any | None]) – Dict input batch following the GeoPrior batch API.
training (bool)
return_maps (bool)
for_train (bool)

Return type:

Important helper#

geoprior.models.subsidence.step_core.physics_core(model, inputs, training, return_maps=False, *, for_train=False)[source]

Compute GeoPrior physics residuals and losses for a batch.

This function implements the shared physics pathway used by both training and evaluation for GeoPrior-style PINN models. It is designed to keep the physics logic consistent across:

train_step() (when physics losses are added to the total loss)
evaluation routines (when physics diagnostics are reported)

At a high level, the function performs:

Input preparation and SI conversions (thickness, head, coords).
Forward pass through the model to obtain data predictions and physics logits.
Mapping of physics logits to bounded physical fields (\(K\), \(S_s\), \(tau\)) and the closure prior \(tau_{phys}\).
Automatic differentiation to obtain PDE derivatives with respect to the model coords.
Chain-rule scaling to SI-consistent derivatives.
Construction of residual maps for: * consolidation relaxation residual, * groundwater flow residual, * time-scale prior residual, * smoothness prior residual, * bounds residual.
Optional nondimensionalization / residual scaling.
Assembly of physics losses, gating schedules, and diagnostic epsilon metrics.

The returned dictionary contains predictions, auxiliary forward outputs, packed physics values (for logging), and optionally the full residual maps and fields.

Parameters:

model (object) –
Model instance providing GeoPrior-style methods and attributes.

The function expects the model to expose (at minimum):
- scaling_kwargsdict
  Resolved scaling and convention payload.
- time_unitsstr or None
  Dataset time unit (for per-second conversions).
- forecast_horizonint
  Horizon length used to tile coords when needed.
- _forward_all(inputs, training=...)callable
  Forward pass returning (y_pred, aux).
- split_data_predictions(x)callable
  Split concatenated data head into subsidence and GWL.
- split_physics_predictions(x)callable
  Split concatenated physics head into (K_logits, Ss_logits, dlogtau_logits, Q_logits).
- pde_modes_activeiterable of str
  Active PDE modes (e.g., {‘consolidation’, ‘gw_flow’}).
- Optional gates: _q_gate(), _subs_resid_gate().
- Optional physics switch: _physics_off().
The function is tolerant to partial capabilities and will short-circuit when physics is disabled, but missing mandatory signals (e.g., thickness) raise errors.
inputs (dict[str, Any | None]) – Dict input batch following the GeoPrior batch API.
training (bool)
return_maps (bool)
for_train (bool)

Return type:

Derivative and residual helpers#

The derivatives module groups the derivative-oriented helpers used by the physics stack. These routines are useful when you want to understand how gradients, rates, or spatial and temporal derivative terms are assembled before they are fed into the residual machinery.

Derivative helpers for GeoPrior PINN blocks.

Goal: keep train_step() and _evaluate_physics_on_batch() consistent and DRY for coordinate chain-rule conversions.

Conventions#

Raw autodiff derivatives are w.r.t. the coordinates tensor fed to call().
This module converts those derivatives to SI-consistent forms: time derivatives to per-second, and spatial derivatives to per-meter (and per-meter squared for second derivatives).

The helper is “conversion-aware”:

If coords are normalized and scaling_kwargs provides coord_ranges_si, those SI spans are used directly (t in seconds, x/y in meters).
Otherwise, it falls back to coord_ranges(), optional deg_to_m(), and finally rate_to_per_second() for time.

It also returns t_range_units_tf (the original time span in time_units) for Q conversion, because Q scaling typically expects the span in the same time units used by the dataset, not seconds.

geoprior.models.subsidence.derivatives.compute_head_pde_derivatives_raw(tape, coords, h_si, K_field, Ss_field)[source]

Compute raw autodiff derivatives for the groundwater-flow PDE.

This helper computes first- and second-order derivatives needed by the GeoPrior groundwater-flow residual using automatic differentiation (AD). All derivatives returned by this function are in the “raw” coordinate units of the coords tensor supplied to the model, without chain-rule rescaling to SI units.

The returned tensors are intended to be passed to ensure_si_derivative_frame() to obtain SI-consistent forms (per-second time derivatives and per-meter spatial derivatives).

Parameters:: tape – Gradient tape that recorded operations connecting h_si, K_field, and Ss_field to coords.

coordsTensor

Coordinate tensor used as the differentiation variable.

Expected shape is (B, H, 3) where the last axis stores coordinates ordered as ['t', 'x', 'y']. The order must be consistent with how dh_dcoords[..., i] is interpreted.

coords may be normalized or unnormalized.
Units may be dataset units or degrees/meters. This function does not apply any unit conversion.

h_siTensor

Hydraulic head in SI-consistent units (or the internal head unit chosen by the pipeline).

Expected shape is (B, H, 1). The tensor must be connected to coords through the computation graph, otherwise AD gradients will be None.

K_fieldTensor

Hydraulic conductivity field \(K\) evaluated on the same batch and horizon grid.

Expected shape is (B, H, 1). The tensor must be connected to coords for spatial gradients to be defined.

Ss_fieldTensor

Specific storage field \(S_s\) evaluated on the same batch and horizon grid.

Expected shape is (B, H, 1). The tensor must be connected to coords for spatial gradients to be defined.

Returns:

grads – Dictionary containing raw derivatives in the coordinate units of coords. Keys include:

'dh_dt_raw': Raw time derivative \(\partial h / \partial t_{raw}\).
'd_K_dh_dx_dx_raw': Raw x-direction divergence term: \(\partial_x (K \partial_x h)\) in raw coord units.
'd_K_dh_dy_dy_raw': Raw y-direction divergence term: \(\partial_y (K \partial_y h)\) in raw coord units.
'dK_dx_raw', 'dK_dy_raw': Raw spatial gradients of \(K\) w.r.t. x and y.
'dSs_dx_raw', 'dSs_dy_raw': Raw spatial gradients of \(S_s\) w.r.t. x and y.

All tensors are expected to have shape (B, H, 1). No scaling by coordinate ranges is applied here.

Return type:

dict of str to Tensor

Raises:

ValueError – If any required gradient is None, indicating the computation graph is not connected to coords or the tape did not watch coords.
ValueError – If any second-order gradients required for the divergence form are None.

Groundwater-flow residual context#

This function provides building blocks for the divergence form used in the groundwater-flow residual:

(100)#\[\begin{split}R_{gw} = S_s \\, \partial_t h - \nabla \cdot (K \\, \nabla h) - Q\end{split}\]

The divergence term in 2D can be expressed as:

(101)#\[\nabla \cdot (K \nabla h) = \partial_x (K \partial_x h) + \partial_y (K \partial_y h)\]

This helper returns the two directional components separately so that downstream code can apply unit conversions and scaling consistently.

Implementation details#

First-order gradients are computed as:

(102)#\[\nabla_{coords} h = \frac{\partial h}{\partial coords}\]

and then split by coordinate axis index.
Second-order divergence terms are computed by differentiating the products K_field * dh_dx_raw and K_field * dh_dy_raw with respect to coords and extracting the x and y components.

Examples

Compute raw derivatives and then convert to SI:

>>> from geoprior.nn.pinn.geoprior.derivatives import (
...     compute_head_pde_derivatives_raw
...     )
>>> with tf.GradientTape(persistent=True) as tape:
...     tape.watch(coords)
...     # forward pass returns h_si, K_field, Ss_field
...     raw = compute_head_pde_derivatives_raw(
...         tape=tape,
...         coords=coords,
...         h_si=h_si,
...         K_field=K_field,
...         Ss_field=Ss_field,
...     )
>>> deriv, meta = ensure_si_derivative_frame(
...     dh_dt_raw=raw["dh_dt_raw"],
...     d_K_dh_dx_dx_raw=raw["d_K_dh_dx_dx_raw"],
...     d_K_dh_dy_dy_raw=raw["d_K_dh_dy_dy_raw"],
...     dK_dx_raw=raw["dK_dx_raw"],
...     dK_dy_raw=raw["dK_dy_raw"],
...     dSs_dx_raw=raw["dSs_dx_raw"],
...     dSs_dy_raw=raw["dSs_dy_raw"],
...     scaling_kwargs=scaling_kwargs,
...     time_units=time_units,
... )

See also

ensure_si_derivative_frame: Convert raw derivatives to SI-consistent derivatives.
geoprior.nn.pinn.geoprior.losses: Physics losses that consume SI-consistent PDE derivatives.

References

Bear, J. Dynamics of Fluids in Porous Media. Dover: Publications, 1988.
Raissi, M., Perdikaris, P., and Karniadakis, G. E.: Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 2019.

geoprior.models.subsidence.derivatives.ensure_si_derivative_frame(*, dh_dt_raw, d_K_dh_dx_dx_raw, d_K_dh_dy_dy_raw, dK_dx_raw, dK_dy_raw, dSs_dx_raw, dSs_dy_raw, scaling_kwargs, time_units, coords_normalized=None, coords_in_degrees=None, eps=1e-12)[source]

Convert autodiff derivative tensors into SI-consistent derivatives.

This helper is the canonical “chain-rule bridge” between raw autodiff gradients taken with respect to the model input coords tensor and the SI-consistent derivatives required by GeoPrior physics losses.

It is designed to keep train_step() and _evaluate_physics_on_batch() consistent and DRY:

Raw derivatives are w.r.t. the coords tensor passed to call().
If coords are normalized, derivatives are rescaled by coordinate spans (and spans squared for second derivatives).
If spatial coords are degrees, spatial derivatives are converted to per-meter forms using a degrees-to-meters factor.
Time derivatives are converted to per-second using time_units unless SI time spans are already supplied.

Parameters:

dh_dt_raw (Tensor) –
Raw autodiff time derivative of head w.r.t. the first coord axis, i.e. \(\partial h / \partial t_{raw}\).

Expected shape is (B, H, 1). The tensor is assumed to be computed w.r.t. the coords tensor fed to call().
d_K_dh_dx_dx_raw (Tensor) –
Raw second-order x-direction PDE term computed as the x component of \(\nabla \cdot (K \nabla h)\) in raw coord units. Conceptually:

(103)\[\partial_x (K \partial_x h)\]

Expected shape is (B, H, 1).
d_K_dh_dy_dy_raw (Tensor) –
Raw second-order y-direction PDE term in raw coord units:

(104)\[\partial_y (K \partial_y h)\]

Expected shape is (B, H, 1).
dK_dx_raw (Tensor) –
Raw spatial gradient of \(K\) in the x direction in raw coord units.

Expected shape is (B, H, 1).
dK_dy_raw (Tensor) –
Raw spatial gradient of \(K\) in the y direction in raw coord units.

Expected shape is (B, H, 1).
dSs_dx_raw (Tensor) –
Raw spatial gradient of \(S_s\) in the x direction in raw coord units.

Expected shape is (B, H, 1).
dSs_dy_raw (Tensor) –
Raw spatial gradient of \(S_s\) in the y direction in raw coord units.

Expected shape is (B, H, 1).
scaling_kwargs –
Scaling and convention payload (resolved config) that describes coordinate normalization and units.

This function primarily consults the following keys:
- coords_normalized. If True, apply span-based chain-rule scaling.
- coord_ranges. Original coordinate spans in dataset units, keyed by 't', 'x', and 'y'. Required when coords_normalized=True.
- coord_ranges_si. Coordinate spans in SI units, keyed by 't', 'x', and 'y' where t is in seconds and x/y are in meters. If present, this is preferred over coord_ranges.
- coords_in_degrees. If True, spatial axes are in degrees and must be converted to meters if SI spans were not already provided.
time_units (str | None)
coords_normalized (bool | None)
coords_in_degrees (bool | None)
eps (float)

Return type:

tuple[dict[str, Tensor], dict[str, Any]]

time_unitsstr or None

Dataset time unit name for the t axis, used to convert the time derivative to per-second when SI time spans are not already provided.

Typical values include 'year', 'day', or 'second'.

coords_normalizedbool, optional

Optional override for scaling_kwargs['coords_normalized']. If provided, this value takes precedence over the payload.

coords_in_degreesbool, optional

Optional override for scaling_kwargs['coords_in_degrees']. If provided, this value takes precedence over the payload.

epsfloat, default 1e-12

Numerical stabilizer added to denominators to avoid division by zero when spans are extremely small or missing.

Returns:

deriv (dict of str to Tensor) – Dictionary of SI-consistent derivative tensors. Keys include:

'dh_dt'
Time derivative converted to per-second: \(\partial h / \partial t\) in SI time.

'd_K_dh_dx_dx' and 'd_K_dh_dy_dy'
Spatial second-derivative PDE terms converted to per-meter squared scaling (via span squared), consistent with the divergence form.

'dK_dx', 'dK_dy', 'dSs_dx', 'dSs_dy'
Spatial gradients converted to per-meter scaling.

The exact physical units of the returned tensors depend on the units of h_si and the representation of K and S_s. The purpose of this function is to enforce correct coordinate scaling (per-second, per-meter, per-meter squared).
meta (dict of str to Any) – Metadata describing which conversion path was used. Important keys include:
- 'used_coord_ranges_si': True if SI spans were taken from coord_ranges_si.
- 'time_already_si': True if an SI time span in seconds was provided.
- 'deg_already_applied': True if x/y spans were already in meters and no degree-to-meter correction was applied.
- 't_range_units_tf': Original time span in dataset time units, retained for downstream Q scaling logic.

Parameters:

dh_dt_raw (Tensor)
d_K_dh_dx_dx_raw (Tensor)
d_K_dh_dy_dy_raw (Tensor)
dK_dx_raw (Tensor)
dK_dy_raw (Tensor)
dSs_dx_raw (Tensor)
dSs_dy_raw (Tensor)
scaling_kwargs (dict[str, Any] | None)
time_units (str | None)
coords_normalized (bool | None)
coords_in_degrees (bool | None)
eps (float)

Return type:

tuple[dict[str, Tensor], dict[str, Any]]

Notes

Chain-rule scaling for normalized coordinates. If normalized coordinates are defined as:

(105)#\[u' = (u - u_0) / \Delta u\]

then derivatives transform as:

(106)#\[\frac{\partial}{\partial u} = \frac{1}{\Delta u} \frac{\partial}{\partial u'}\]

and second derivatives as:

(107)#\[\frac{\partial^2}{\partial u^2} = \frac{1}{(\Delta u)^2} \frac{\partial^2}{\partial (u')^2}\]

This function applies these rules using either coord_ranges_si (preferred) or coord_ranges plus unit conversion.

Degrees to meters conversion. If spatial coords are degrees (longitude/latitude), the function converts spatial derivative scaling using a degrees-to-meters factor derived from the scaling payload. This is only applied when SI spans were not already provided.

Examples

Convert derivatives for normalized coords with SI spans:

>>> from geoprior.nn.pinn.geoprior.derivatives import (
...    ensure_si_derivative_frame
...    )
>>> deriv, meta = ensure_si_derivative_frame(
...     dh_dt_raw=dh_dt_raw,
...     d_K_dh_dx_dx_raw=dKdhx_dx_raw,
...     d_K_dh_dy_dy_raw=dKdhy_dy_raw,
...     dK_dx_raw=dK_dx_raw,
...     dK_dy_raw=dK_dy_raw,
...     dSs_dx_raw=dSs_dx_raw,
...     dSs_dy_raw=dSs_dy_raw,
...     scaling_kwargs={
...         "coords_normalized": True,
...         "coord_ranges": {"t": 7.0, "x": 4.4e4, "y": 3.9e4},
...         "coord_ranges_si": {
...             "t": 2.2e8, "x": 4.4e4, "y": 3.9e4
...         },
...     },
...     time_units="year",
... )
>>> bool(meta["used_coord_ranges_si"])
True

Fallback when SI spans are absent (time converted using time_units):

>>> deriv, meta = ensure_si_derivative_frame(
...     dh_dt_raw=dh_dt_raw,
...     d_K_dh_dx_dx_raw=dKdhx_dx_raw,
...     d_K_dh_dy_dy_raw=dKdhy_dy_raw,
...     dK_dx_raw=dK_dx_raw,
...     dK_dy_raw=dK_dy_raw,
...     dSs_dx_raw=dSs_dx_raw,
...     dSs_dy_raw=dSs_dy_raw,
...     scaling_kwargs={
...         "coords_normalized": True,
...         "coord_ranges": {"t": 7.0, "x": 4.4e4, "y": 3.9e4},
...     },
...     time_units="year",
... )
>>> bool(meta["time_already_si"])
False

See also

compute_head_pde_derivatives_raw: Compute raw autodiff derivatives w.r.t. input coords.
geoprior.nn.pinn.geoprior.maths.rate_to_per_second: Convert a time rate from dataset units to per-second.
geoprior.nn.pinn.geoprior.utils.coord_ranges: Extract coordinate spans from a scaling payload.
geoprior.nn.pinn.geoprior.utils.deg_to_m: Convert degrees to meters scaling for spatial axes.

References

Raissi, M., Perdikaris, P., and Karniadakis, G. E.
Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 2019.

Batch and transport helpers#

The batch_io module provides helpers for moving structured batches into the forms expected by the subsidence models and their staged workflow. It is especially useful when tracing how batch dictionaries or arrays are normalized before they enter the shared physics core.

Batch.io

geoprior.models.subsidence.batch_io.select_q(pred, quantiles=None, q=0.5, fallback='mean')[source]

Select q-quantile from pred.

pred:

(B,H,Q,O) -> returns (B,H,O)
(B,H,Q,1) -> returns (B,H,1)
otherwise returned as-is.

geoprior.models.subsidence.batch_io.tile_true_to_quantiles(y_true, y_pred)[source]

Make y_true compatible with y_pred when y_pred has quantile axis.

y_true: (B,H,O) and y_pred: (B,H,Q,O) -> return y_true_q: (B,H,Q,O) Else return y_true unchanged.

Debugging and scientific inspection#

The debugs module contains helpers intended for inspection-oriented workflows. These routines are useful when you need to verify scientific assumptions, track internal state, or make sense of unexpected training or evaluation behavior.

Debug helpers for GeoPriorSubsNet.

Keep all verbosity + shape/unit printing here so _geoprior_subnet.py stays clean.

All functions are safe to call inside tf.function: they use tf.print and TensorFlow assertions. - Compare in-memory vs loaded inference model on the same batch. - Report prediction diffs + weight name/shape diffs + key attribute digests.

geoprior.models.subsidence.debugs.weight_diff_report(m1, m2, *, top=30, include_ok=False, include_extra=False)[source]

Compare weights between two models using stable keys.

Each row is:: (max_abs_diff or inf, tag, weight_id, shape_info)
Where tag in:: {“OK”, “MISSING”, “SHAPE”, “EXTRA”}

Notes

Uses w.path when available (best).
Otherwise uses name + occurrence index.
Set top<=0 to return all rows.

Parameters:

m1 (Any)
m2 (Any)
top (int)
include_ok (bool)
include_extra (bool)

Return type:

list[tuple[float, str, str, Any]]

geoprior.models.subsidence.debugs.model_scaling_digest(model)[source]

Stable digest of model.scaling_kwargs to verify the same config survived reload.

Parameters:: model (Any)
Return type:: str

geoprior.models.subsidence.debugs.debug_model_reload(mem_model, load_model, dataset, *, pred_key='subs_pred', also_check=None, top_weights=30, atol=1e-06, rtol=1e-06, log_fn=None)[source]

Run a compact reload debug on one batch and return a dict report.

Compares predictions (max/mean abs diff) for pred_key (+ optional keys).
Compares weights by name (MISSING/SHAPE/OK).
Compares scaling_kwargs digest + time_units attribute.

Parameters:

mem_model (Any)
load_model (Any)
dataset (Iterable)
pred_key (str)
also_check (list[str] | None)
top_weights (int)
atol (float)
rtol (float)
log_fn (Callable[[str], None] | None)

Return type:

geoprior.models.subsidence.debugs.dbg_on(verbose, level)[source]

Return True if verbose is strictly above level.

Parameters:

verbose (int)
level (int)

Return type:

bool

geoprior.models.subsidence.debugs.dbg_run_first_iter(*, verbose, level, iterations, fn)[source]

Run fn() only at optimizer.iterations == 0.

Graph-safe: uses tf.cond and returns a dummy scalar.

Parameters:

verbose (int)
level (int)
iterations (Tensor)

Return type:

None

geoprior.models.subsidence.debugs.dbg_stats(tag, x)[source]

Print min/max/mean (graph-safe).

Parameters:

tag (str)
x (Tensor)

Return type:

None

geoprior.models.subsidence.debugs.dbg_pde_divergence_maxabs(*, verbose, raw_dKdhx_dcoords, raw_d_K_dh_dx_dx, raw_d_K_dh_dy_dy, d_K_dh_dx_dx=None, d_K_dh_dy_dy=None, level=7, prefix='pde/div')[source]

Print max-abs diagnostics for the divergence terms, before and optionally after normalization/chain-rule correction.

Parameters:

verbose (int)
raw_dKdhx_dcoords (Tensor)
raw_d_K_dh_dx_dx (Tensor)
raw_d_K_dh_dy_dy (Tensor)
d_K_dh_dx_dx (Any | None)
d_K_dh_dy_dy (Any | None)
level (int)
prefix (str)

Return type:

None

geoprior.models.subsidence.debugs.dbg_gw_units_and_sec_scale(*, verbose, gw_units, gw_res_before, gw_res_after, level=7, prefix='gw/units')[source]

Print GW residual diagnostics before/after applying sec_u scaling.

Call this right after you do:: gw_res_before = gw_res gw_res = gw_res * sec_u gw_res_after = gw_res

Parameters:

gw_units (Any) – Usually resolve_gw_units(sk). Keep it as Any so callers can pass python strings without TF ops.
verbose (int)
gw_res_before (Tensor)
gw_res_after (Tensor)
level (int)
prefix (str)

Return type:

None

Notes

We print RMS to catch accidental unit explosions.

geoprior.models.subsidence.debugs.dbg_mae(tag, y, yhat)[source]

Print batch MAE for y vs yhat.

Parameters:

tag (str)
y (Tensor)
yhat (Tensor)

Return type:

None

geoprior.models.subsidence.debugs.dbg_chk_finite(tag, x)[source]

Assert finite, return x (small helper).

Parameters:

tag (str)
x (Tensor)

Return type:

Tensor

geoprior.models.subsidence.debugs.dbg_step0_inputs_targets(*, verbose, inputs, targets, level=12)[source]

Parameters:

verbose (int)
inputs (dict[str, Any])
targets (Any)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step1_thickness(*, verbose, H_field, H_si, level=12)[source]

Parameters:

verbose (int)
H_field (Tensor)
H_si (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step2_coords_checks(*, verbose, coords, inputs, level=12)[source]

Parameters:

verbose (int)
coords (Tensor)
inputs (dict[str, Any])
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_units_once(*, verbose, iterations, targets, gwl_pred_final, s_pred_final, quantiles, level=7)[source]

Parameters:

verbose (int)
iterations (Tensor)
targets (dict[str, Tensor])
gwl_pred_final (Tensor)
s_pred_final (Tensor)
quantiles (Sequence[float] | None)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_assert_data_layout(*, verbose, data_final, data_mean_raw, quantiles, level=12)[source]

Parameters:

verbose (int)
data_final (Tensor)
data_mean_raw (Any | None)
quantiles (Sequence[float] | None)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step3_mean_head(*, verbose, gwl_mean_raw, gwl_si, h_si, level=12)[source]

Parameters:

verbose (int)
gwl_mean_raw (Tensor)
gwl_si (Tensor)
h_si (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step31_forward_outputs(*, verbose, data_final, s_pred_final, gwl_pred_final, data_mean_raw, phys_mean_raw, level=12)[source]

Parameters:

verbose (int)
data_final (Tensor)
s_pred_final (Tensor)
gwl_pred_final (Tensor)
data_mean_raw (Any | None)
phys_mean_raw (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step33_physics_logits(*, verbose, K_logits, Ss_logits, dlogtau_logits, Q_logits, K_base, Ss_base, dlogtau_base, level=12)[source]

Parameters:

verbose (int)
K_logits (Tensor)
Ss_logits (Tensor)
dlogtau_logits (Tensor)
Q_logits (Any | None)
K_base (Tensor)
Ss_base (Tensor)
dlogtau_base (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step33_physics_fields(*, verbose, K_field, Ss_field, tau_field, tau_phys, Hd_eff, delta_log_tau, logK, logSs, log_tau, log_tau_phys, level=12)[source]

Parameters:

verbose (int)
K_field (Tensor)
Ss_field (Tensor)
tau_field (Tensor)
tau_phys (Tensor)
Hd_eff (Tensor)
delta_log_tau (Tensor)
logK (Tensor)
logSs (Tensor)
log_tau (Tensor)
log_tau_phys (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step4_ad_raw(*, verbose, dh_dcoords, dh_dt_raw, dh_dx_raw, dh_dy_raw, K_dh_dx, K_dh_dy, dKdhx_dcoords, dKdhy_dcoords, d_K_dh_dx_dx_raw, d_K_dh_dy_dy_raw, dK_dcoords, dSs_dcoords, dK_dx_raw, dK_dy_raw, dSs_dx_raw, dSs_dy_raw, level=12)[source]

Parameters:

verbose (int)
dh_dcoords (Tensor)
dh_dt_raw (Tensor)
dh_dx_raw (Tensor)
dh_dy_raw (Tensor)
K_dh_dx (Tensor)
K_dh_dy (Tensor)
dKdhx_dcoords (Tensor)
dKdhy_dcoords (Tensor)
d_K_dh_dx_dx_raw (Tensor)
d_K_dh_dy_dy_raw (Tensor)
dK_dcoords (Tensor)
dSs_dcoords (Tensor)
dK_dx_raw (Tensor)
dK_dy_raw (Tensor)
dSs_dx_raw (Tensor)
dSs_dy_raw (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step41_si_grads(*, verbose, dh_dt, d_K_dh_dx_dx, d_K_dh_dy_dy, dK_dx, dK_dy, dSs_dx, dSs_dy, level=12)[source]

Parameters:

verbose (int)
dh_dt (Tensor)
d_K_dh_dx_dx (Tensor)
d_K_dh_dy_dy (Tensor)
dK_dx (Tensor)
dK_dy (Tensor)
dSs_dx (Tensor)
dSs_dy (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step5_q_source(*, verbose, Q_si, dh_dt, level=12)[source]

Parameters:

verbose (int)
Q_si (Tensor)
dh_dt (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_cons_units_rms(*, verbose, sk, cons_res, level=7)[source]

Parameters:

verbose (int)
sk (dict[str, Any])
cons_res (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step6_consolidation(*, verbose, allow_resid, cons_active, s_mean_raw, s_pred_si, dt_units, s0_cum_11, s_inc_pred, s_state, h_ref_si_11, h_state, cons_step_m, cons_res)[source]

Parameters:

verbose (int)
allow_resid (bool)
cons_active (bool)
s_mean_raw (Tensor)
s_pred_si (Tensor)
dt_units (Tensor)
s0_cum_11 (Tensor)
s_inc_pred (Tensor)
s_state (Tensor)
h_ref_si_11 (Tensor)
h_state (Tensor)
cons_step_m (Tensor)
cons_res (Tensor)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step7_residuals(*, verbose, gw_res, prior_res, smooth_res, loss_mv, bounds_res, loss_bounds, level=12)[source]

Parameters:

verbose (int)
gw_res (Tensor)
prior_res (Tensor)
smooth_res (Tensor)
loss_mv (Tensor)
bounds_res (Tensor)
loss_bounds (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step8_scaling(*, verbose, cons_res_raw, gw_res_raw, cons_res, gw_res, level=7)[source]

Parameters:

verbose (int)
cons_res_raw (Tensor)
gw_res_raw (Tensor)
cons_res (Tensor)
gw_res (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_chk_scales(*, verbose, scales, level=2)[source]

Parameters:

verbose (int)
scales (dict[str, Tensor])
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_chk_core_finite(*, verbose, cons_res, gw_res, tau_field, K_field, Ss_field, level=2)[source]

Parameters:

verbose (int)
cons_res (Tensor)
gw_res (Tensor)
tau_field (Tensor)
K_field (Tensor)
Ss_field (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step9_losses(*, verbose, data_loss=None, loss_cons=None, loss_gw=None, loss_prior=None, loss_smooth=None, physics_loss_raw=None, physics_loss_scaled=None, total_loss=None, level=7)[source]

Debug-print loss scalars (only those provided).

Parameters:

verbose (int)
data_loss (Any | None)
loss_cons (Any | None)
loss_gw (Any | None)
loss_prior (Any | None)
loss_smooth (Any | None)
physics_loss_raw (Any | None)
physics_loss_scaled (Any | None)
total_loss (Any | None)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step10_grads(*, verbose, trainable_vars, grads, level=9)[source]

Parameters:

verbose (int)
trainable_vars (Sequence[Tensor])
grads (Sequence[Any | None])
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_term_grads_finite(*, verbose, debug_grads, trainable_vars, data_loss, terms_scaled, tape, level=1)[source]

Parameters:

verbose (int)
debug_grads (bool)
trainable_vars (Sequence[Tensor])
data_loss (Tensor)
terms_scaled (dict[str, Tensor])
tape (Any)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_done_apply_gradients(*, debug_grads=False, verbose=1)[source]

Parameters:

debug_grads (bool)
verbose (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_select_q(y, quantiles, *, q=0.5)[source]

Select a quantile slice if y is (B,H,Q,1)/(B,H,Q).

If quantiles is None, returns y as-is.

Parameters:

y (Tensor)
quantiles (Sequence[float] | None)
q (float)

Return type:

Tensor

geoprior.models.subsidence.debugs.dbg_step5_q(*, verbose, Q_si, dh_dt, level=12)[source]

Print Q source term block.

Parameters:

verbose (int)
Q_si (Tensor)
dh_dt (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step8_residual_scale_stats(*, verbose, level=3, cons_res_raw, cons_scale, gw_res_raw, gw_scale)[source]

Debug block for residuals + scaling stats.

Replaces:: _stats(“cons_res_raw”, cons_res) _stats(“cons_scale”, scales[“cons_scale”]) _stats(“gw_res_raw”, gw_res) _stats(“gw_scale”, scales[“gw_scale”])

Parameters:

verbose (int)
level (int)
cons_res_raw (Tensor)
cons_scale (Tensor)
gw_res_raw (Tensor)
gw_scale (Tensor)

Return type:

None

geoprior.models.subsidence.debugs.dbg_dt_debug(*, verbose, level=3, time_units, dt_units, t)[source]

Debug dt conversion and t-grid sanity.

Replaces the “dt debug” block.

Parameters:

verbose (int)
level (int)
time_units (str)
dt_units (Tensor)
t (Tensor)

Return type:

None

geoprior.models.subsidence.debugs.dbg_call_nonfinite(*, verbose, level=9, coords_for_decoder, H_si, K_base, Ss_base, dlogtau_base, tau_field)[source]

Debug non-finite checks for call() internal tensors.

Replaces:: tf_print_nonfinite(“call/coords_for_decoder”, coords_for_decoder) …

Parameters:

verbose (int)
level (int)
coords_for_decoder (Tensor)
H_si (Tensor)
K_base (Tensor)
Ss_base (Tensor)
dlogtau_base (Tensor)
tau_field (Tensor)

Return type:

None

geoprior.models.subsidence.debugs.dbg_step3_residual_scales(*, verbose, cons_res, gw_res, scales, level=3)[source]

Print raw residual stats + scaling factors.

Parameters:

verbose (int)
cons_res (Tensor)
gw_res (Tensor)
scales (dict)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_dt_diag(*, verbose, time_units, dt_units, t, level=3)[source]

Print dt consistency checks in time_units and seconds.

Parameters:

verbose (int)
time_units (str)
dt_units (Tensor)
t (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_call_nonfinite_diag(*, verbose, coords_for_decoder, H_si, K_base, Ss_base, dlogtau_base, tau_field, level=9)[source]

Print non-finite diagnostics inside call().

Parameters:

verbose (int)
coords_for_decoder (Tensor)
H_si (Tensor)
K_base (Tensor)
Ss_base (Tensor)
dlogtau_base (Tensor)
tau_field (Tensor)
level (int)

Return type:

None

geoprior.models.subsidence.debugs.dbg_gw_grad_flux_rms(*, verbose, dh_dx_raw, dh_dy_raw, K_field, level=3, prefix='gw/gradflux')[source]

Print RMS diagnostics for spatial head gradients and Darcy-like flux terms K*∂h/∂x, K*∂h/∂y (raw coord units).

Replaces:: tf_print(“to_rms(dh_dx)=”, to_rms(dh_dx_raw)) tf_print(“to_rms(dh_dy)=”, to_rms(dh_dy_raw)) tf_print(“to_rms(K_field * dh_dx)=”, to_rms(K_field * dh_dx_raw)) tf_print(“to_rms(K_field * dh_dy)=”, to_rms(K_field * dh_dy_raw))

Parameters:

verbose (int)
dh_dx_raw (Tensor)
dh_dy_raw (Tensor)
K_field (Tensor)
level (int)
prefix (str)

Return type:

None

Log-offset diagnostics#

The log_offsets_diagnostics module focuses on diagnostics for inferred log-offset fields and related quantities. It is particularly relevant when you are studying prior anchoring, offset magnitudes, or identifiability behavior.

Diagnostics for subsidence log-offset policies and payloads.

geoprior.models.subsidence.log_offsets_diagnostics.run_sm3_offsets_from_payload(physics_npz_path, outdir=None, city=None, model_name='GeoPriorSubsNet')[source]

High-level driver: compute SM3 diagnostics from a physics payload.

Parameters:

physics_npz_path (str) – Path to *_phys_payload_run_val.npz as written by GeoPriorSubsNet.export_physics_payload().
outdir (str, optional) – Directory where CSVs and plots are written. If None, defaults to the directory of physics_npz_path.
city (str, optional) – City name for filenames.
model_name (str, default "GeoPriorSubsNet") – Model name for filenames.

Returns:

result – Dictionary with keys: - ‘raw_csv’ - ‘summary_csv’ - ‘plots’ (list of paths)

Return type: