fusionlab.nn.pinn.PiTGWFlow

class fusionlab.nn.pinn.PiTGWFlow

Bases: Model, NNLearner

Physics-Informed Transient Groundwater Flow.

A self-contained Physics-Informed Neural Network (PINN) designed to solve the 2D/3D transient groundwater flow equation. This model leverages a simple Multi-Layer Perceptron (MLP) to approximate the hydraulic head \(h\) as a continuous function of time and space, \(h = h(t, x, y)\).

The model is trained by minimizing the residual of the governing PDE at a set of collocation points, rather than by fitting to observed data. The core of the model is its custom training and evaluation steps, which compute the necessary spatial and temporal derivatives to enforce the physical law.

The governing equation solved by this PINN is:

\[S_s \frac{\partial h}{\partial t} - K \left( \frac{\partial^2 h}{\partial x^2} + \frac{\partial^2 h}{\partial y^2} \right) - Q = 0\]

Where \(K\) is the hydraulic conductivity, \(S_s\) is the specific storage, and \(Q\) is a source/sink term. These physical parameters can be set as fixed constants or as trainable variables.

Parameters:

• hidden_units (int or list of int, optional) –

  Defines the architecture of the internal MLP. - If an int, three hidden layers of that size are created. - If a list, its first three values are used as the sizes

  for the three hidden layers.

  Defaults to [32, 32, 32].

• activation (str, default 'tanh') – The activation function to use in the hidden layers of the MLP. ‘tanh’ is often recommended for PINNs.

• learning_rate (float, default 1e-3) – The learning rate for the Adam optimizer used to train the network weights and any learnable physical parameters.

• K (float or LearnableK, default 1.0) – The hydraulic conductivity. Can be provided as a fixed Python float or as a LearnableK instance to be optimized during training.

• Ss (float or LearnableSs, default 1e-4) – The specific storage. Can be provided as a fixed Python float or as a LearnableSs instance to be optimized during training.

• Q (float or LearnableQ, default 0.0) – The volumetric source/sink term. Can be provided as a fixed Python float or as a LearnableQ instance to be optimized during training.

• gw_coeffs_config (dict, optional) –

  A dictionary to conveniently configure physical parameters. If provided, its values will supersede the individual K, Ss, and Q arguments. For example:

  gw_coeffs_config = {'K': LearnableK(0.5), 'Ss': 1e-5}
  

• name (str, default 'PiTGWFlow') – The name of the Keras model.

• **kwargs – Additional keyword arguments passed to the parent tf.keras.Model constructor.

Notes

This model is a pure physics solver, meaning its loss function is derived entirely from the PDE residual. It does not perform data fitting in the traditional sense. Consequently, the y_true component of a dataset provided to fit() or evaluate() is ignored.

The model implements custom train_step and test_step methods. Within these steps, tf.GradientTape is used to calculate the first and second-order derivatives of the network’s output (\(h\)) with respect to its inputs (\(t, x, y\)). This is necessary to compute the PDE residual, which forms the loss.

When a physical parameter is made learnable (e.g., by passing K=LearnableK(...)), it is automatically added to the model’s list of trainable variables and is updated via backpropagation to minimize the physics loss.

See also

fusionlab.nn.pinn.TransFlowSubsNet

A more complex, hybrid data-physics model that couples groundwater flow with land subsidence and is trained on both physical laws and observed data.

Examples

>>> import tensorflow as tf
>>> from fusionlab.nn.pinn import PiTGWFlow
>>> from fusionlab.params import LearnableK, LearnableSs
...
>>> # 1. Instantiate the model with a learnable K and a fixed Ss
>>> model = PiTGWFlow(
...     hidden_units=[64, 64, 64],
...     K=LearnableK(initial_value=0.5),
...     Ss=1e-5,
...     Q=0.0
... )
...
>>> # 2. Create a dataset of collocation points (t, x, y)
>>> n_points = 1000
>>> coords = {
...     "t": tf.random.uniform((n_points, 1)),
...     "x": tf.random.uniform((n_points, 1)),
...     "y": tf.random.uniform((n_points, 1)),
... }
>>> # Dummy targets are needed for the Keras API, but are ignored
>>> dummy_y = tf.zeros((n_points, 1))
>>> dataset = tf.data.Dataset.from_tensor_slices((coords, dummy_y)).batch(32)
...
>>> # 3. Compile and train the model
>>> # The optimizer is taken from the model instance.
>>> model.compile()
>>> history = model.fit(dataset, epochs=10)
...
>>> # The history will contain the physics-based loss
>>> print(list(history.history.keys()))
['pde_loss']

__init__(hidden_units=None, activation='tanh', learning_rate=0.001, K=1.0, Ss=0.0001, Q=0.0, gw_coeffs_config=None, name='PiTGWFlow', **kwargs)

Parameters:

• hidden_units (int | List[int] | None)

• activation (str)

• learning_rate (float)

• K (float | LearnableK)

• Ss (float | LearnableSs)

• Q (float | LearnableQ)

• gw_coeffs_config (Dict[str, Any])

• name (str)

Methods

__init__([hidden_units, activation, ...])
add_loss(loss)	Can be called inside of the call() method to add a scalar loss.
add_metric(*args, **kwargs)
add_variable(shape, initializer[, dtype, ...])	Add a weight variable to the layer.
add_weight([shape, initializer, dtype, ...])	Add a weight variable to the layer.
build(input_shape)
build_from_config(config)	Builds the layer's states with the supplied config dict.
call(inputs[, training])	Performs the forward pass of the network.
compile([optimizer, loss, loss_weights, ...])	Configures the model for training.
compile_from_config(config)	Compiles the model with the information given in config.
compiled_loss(y, y_pred[, sample_weight, ...])
compute_loss([x, y, y_pred, sample_weight, ...])	Compute the total loss, validate it, and return it.
compute_mask(inputs, previous_mask)
compute_metrics(x, y, y_pred[, sample_weight])	Update metric states and collect all metrics to be returned.
compute_output_shape(*args, **kwargs)
compute_output_spec(*args, **kwargs)
count_params()	Count the total number of scalars composing the weights.
evaluate([x, y, batch_size, verbose, ...])	Returns the loss value & metrics values for the model in test mode.
export(filepath[, format, verbose, ...])	Export the model as an artifact for inference.
fit([x, y, batch_size, epochs, verbose, ...])	Trains the model for a fixed number of epochs (dataset iterations).
from_config(config[, custom_objects])	Creates a model instance from its configuration dictionary.
get_build_config()	Returns a dictionary with the layer's input shape.
get_compile_config()	Returns a serialized config with information for compiling the model.
get_config()	Returns the serializable configuration of the model.
get_layer([name, index])	Retrieves a layer based on either its name (unique) or index.
get_metrics_result()	Returns the model's metrics values as a dict.
get_params([deep])	Get the parameters for this learner.
get_state_tree([value_format])	Retrieves tree-like structure of model variables.
get_weights()	Return the values of layer.weights as a list of NumPy arrays.
help(**kwargs)
load(file_path[, format])	Load the learner's state from a specified file in the desired format.
load_own_variables(store)	Loads the state of the layer.
load_weights(filepath[, skip_mismatch])	Load the weights from a single file or sharded files.
loss(y, y_pred[, sample_weight])
make_predict_function([force])
make_test_function([force])
make_train_function([force])
predict(x[, batch_size, verbose, steps, ...])	Generates output predictions for the input samples.
predict_on_batch(x)	Returns predictions for a single batch of samples.
predict_step(data)
quantize(mode[, config])	Quantize the weights of the model.
quantized_build(input_shape, mode)
quantized_call(*args, **kwargs)
rematerialized_call(layer_call, *args, **kwargs)	Enable rematerialization dynamically for layer's call method.
reset_metrics()
save(filepath[, overwrite, zipped])	Saves a model as a .keras file.
save_own_variables(store)	Saves the state of the layer.
save_weights(filepath[, overwrite, ...])	Saves all weights to a single file or sharded files.
set_params(**params)	Set the parameters of this learner.
set_state_tree(state_tree)	Assigns values to variables of the model.
set_weights(weights)	Sets the values of layer.weights from a list of NumPy arrays.
stateless_call(trainable_variables, ...[, ...])	Call the layer without any side effects.
stateless_compute_loss(trainable_variables, ...)
summary([line_length, positions, print_fn, ...])	Prints a string summary of the network.
symbolic_call(*args, **kwargs)
test_on_batch(x[, y, sample_weight, return_dict])	Test the model on a single batch of samples.
test_step(data)	Defines the logic for a single evaluation step.
to_json(**kwargs)	Returns a JSON string containing the network configuration.
train_on_batch(x[, y, sample_weight, ...])	Runs a single gradient update on a single batch of data.
train_step(data)	Defines the logic for a single optimization step.

Attributes

compiled_metrics
compute_dtype	The dtype of the computations performed by the layer.
distribute_reduction_method
distribute_strategy
dtype	Alias of layer.variable_dtype.
dtype_policy
input	Retrieves the input tensor(s) of a symbolic operation.
input_dtype	The dtype layer inputs should be converted to.
input_spec
jit_compile
layers
losses	List of scalar losses from add_loss, regularizers and sublayers.
metrics	List of all metrics.
metrics_names
metrics_variables	List of all metric variables.
my_params
non_trainable_variables	List of all non-trainable layer state.
non_trainable_weights	List of all non-trainable weight variables of the layer.
output	Retrieves the output tensor(s) of a layer.
path	The path of the layer.
quantization_mode	The quantization mode of this layer, None if not quantized.
run_eagerly
supports_masking	Whether this layer supports computing a mask using compute_mask.
trainable	Settable boolean, whether this layer should be trainable or not.
trainable_variables	List of all trainable layer state.
trainable_weights	List of all trainable weight variables of the layer.
variable_dtype	The dtype of the state (weights) of the layer.
variables	List of all layer state, including random seeds.
weights	List of all weight variables of the layer.

__init__(hidden_units=None, activation='tanh', learning_rate=0.001, K=1.0, Ss=0.0001, Q=0.0, gw_coeffs_config=None, name='PiTGWFlow', **kwargs)

Parameters:

• hidden_units (int | List[int] | None)

• activation (str)

• learning_rate (float)

• K (float | LearnableK)

• Ss (float | LearnableSs)

• Q (float | LearnableQ)

• gw_coeffs_config (Dict[str, Any])

• name (str)

call(inputs, training=False)

Performs the forward pass of the network.

This method maps the input spatio-temporal coordinates \((t, x, y)\) to the predicted hydraulic head \(h\). It is designed to be flexible, accepting coordinates in several formats.

Parameters:

• inputs (tf.Tensor or dict) –

  The input coordinates. Supported formats include: - A dictionary with keys 't', 'x', and 'y',

  where each value is a tensor of shape (batch, time_steps, 1).

  • A single concatenated tensor of shape (batch, time_steps, 3).

• training (bool, optional) – Indicates whether the model is in training mode. This is passed to internal Keras layers. Defaults to False.

Returns:

A tensor of shape (batch, time_steps, 1) containing the predicted hydraulic head \(h\) at each input coordinate.

Return type:

tf.Tensor

train_step(data)

Defines the logic for a single optimization step.

This method overrides the default Keras training behavior to implement the physics-informed loss. In each step, it computes the residual of the governing PDE and uses it as the loss to update all trainable variables, including both the network weights and any learnable physical parameters.

The process involves: 1. Computing derivatives of the output \(h\) with respect

to inputs \((t, x, y)\) using tf.GradientTape.

2. Assembling the PDE residual using these derivatives.

3. Calculating the mean squared error of the residual.

4. Applying gradients to update the model’s variables.

Parameters:

data (tuple) – A tuple of (inputs, targets) as provided by a tf.data.Dataset. The inputs are a dictionary of coordinate tensors, while the targets are ignored as the loss is unsupervised.

Returns:

A dictionary mapping the metric name 'pde_loss' to its scalar value for this training step.

Return type:

dict

test_step(data)

Defines the logic for a single evaluation step.

This method overrides the default Keras evaluation behavior to report the model’s physics fidelity. It calculates the PDE residual on the validation or test data but does not perform any weight updates. This is crucial for assessing how well the model adheres to the physical laws on unseen collocation points.

Parameters:

data (tuple) – A tuple of (inputs, targets) from the validation or test dataset. The inputs are a dictionary of coordinate tensors; the targets are ignored.

Returns:

A dictionary mapping the metric name 'pde_loss' to its scalar value for the evaluation step.

Return type:

dict

get_config()

Returns the serializable configuration of the model.

This method allows the model to be saved and re-instantiated without losing its architectural and parameter setup. It captures all initialization arguments, including the configuration of the MLP and the physical parameters.

Returns:

A dictionary containing all constructor parameters needed to recreate the model instance. Learnable parameter objects (e.g., LearnableK) are serialized as part of this config.

Return type:

dict

classmethod from_config(config, custom_objects=None)

Creates a model instance from its configuration dictionary.

This class method is the counterpart to get_config and is used by Keras’s loading utilities to reconstruct a model from its saved configuration.

Parameters:

• config (dict) – The configuration dictionary, typically generated by get_config().

• custom_objects (dict, optional) – A dictionary mapping names to custom classes or functions. Keras uses this to deserialize custom objects like LearnableK.

Returns:

A new instance of the PiTGWFlow model.

Return type:

PiTGWFlow

help(**kwargs)

my_params = PiTGWFlow(     hidden_units=None,     activation='tanh',     learning_rate=0.001,     K=1.0,     Ss=0.0001,     Q=0.0,     gw_coeffs_config=None,     name='PiTGWFlow' )