geoprior.models.subsidence.models#

Subsidence PINN models

Classes

GeoPriorSubsNet(*args, **kwargs)

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

PoroElasticSubsNet(*args, **kwargs)

Poroelastic surrogate variant of GeoPriorSubsNet.
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:
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:
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:

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

dict

Notes

Step outline. This training step performs the following stages:

  1. 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.

  2. 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.

  3. 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.

  4. 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.

  5. 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.

  6. 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.

  7. 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:

dict

Notes

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

  1. Unpack and canonicalize targets

    Targets are normalized into a stable dict structure and reordered by output key contract.

  2. 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.

  3. 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.

  4. 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.

  5. 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.

  6. 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:

GeoPriorSubsNet

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:

  1. Build a custom_objects registry that includes all GeoPrior wrappers and scaling configuration classes needed for safe deserialization.

  2. Rehydrate wrapper objects stored as Keras-serialized dicts ({"class_name": ..., "config": ...}) for keys such as mv, kappa, gamma_w, and h_ref.

  3. 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:

GeoPriorSubsNet

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: