geoprior.models.subsidence.derivatives#

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.

Functions

`compute_head_pde_derivatives_raw`(tape, ...)	Compute raw autodiff derivatives for the groundwater-flow PDE.
`ensure_si_derivative_frame`(*, dh_dt_raw, ...)	Convert autodiff derivative tensors into SI-consistent derivatives.

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:

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

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

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

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

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

(6)#\[u' = (u - u_0) / \Delta u\]

then derivatives transform as:

(7)#\[\frac{\partial}{\partial u} = \frac{1}{\Delta u} \frac{\partial}{\partial u'}\]

and second derivatives as:

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