fusionlab.nn.pinn.PiTGWFlow

class fusionlab.nn.pinn.PiTGWFlow[source]

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)[source]
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(losses, **kwargs)

Add loss tensor(s), potentially dependent on layer inputs.

add_metric(value[, name])

Adds metric tensor to the layer.

add_update(updates)

Add update op(s), potentially dependent on layer inputs.

add_variable(*args, **kwargs)

Deprecated, do NOT use! Alias for add_weight.

add_weight([name, shape, dtype, ...])

Adds a new variable to the layer.

build(input_shape)

Builds the model based on input shapes received.

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, metrics, ...])

Configures the model for training.

compile_from_config(config)

Compiles the model with the information given in config.

compute_loss([x, y, y_pred, sample_weight])

Compute the total loss, validate it, and return it.

compute_mask(inputs[, mask])

Computes an output mask tensor.

compute_metrics(x, y, y_pred, sample_weight)

Update metric states and collect all metrics to be returned.

compute_output_shape(input_shape)

Computes the output shape of the layer.

compute_output_signature(input_signature)

Compute the output tensor signature of the layer based on the inputs.

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.

evaluate_generator(generator[, steps, ...])

Evaluates the model on a data generator.

export(filepath)

Create a SavedModel artifact for inference (e.g. via TF-Serving).

finalize_state()

Finalizes the layers state after updating layer weights.

fit([x, y, batch_size, epochs, verbose, ...])

Trains the model for a fixed number of epochs (dataset iterations).

fit_generator(generator[, steps_per_epoch, ...])

Fits the model on data yielded batch-by-batch by a Python generator.

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_input_at(node_index)

Retrieves the input tensor(s) of a layer at a given node.

get_input_mask_at(node_index)

Retrieves the input mask tensor(s) of a layer at a given node.

get_input_shape_at(node_index)

Retrieves the input shape(s) of a layer at a given node.

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_output_at(node_index)

Retrieves the output tensor(s) of a layer at a given node.

get_output_mask_at(node_index)

Retrieves the output mask tensor(s) of a layer at a given node.

get_output_shape_at(node_index)

Retrieves the output shape(s) of a layer at a given node.

get_params([deep])

Get the parameters for this learner.

get_weight_paths()

Retrieve all the variables and their paths for the model.

get_weights()

Retrieves the weights of the model.

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, ...])

Loads all layer weights from a saved files.

make_predict_function([force])

Creates a function that executes one step of inference.

make_test_function([force])

Creates a function that executes one step of evaluation.

make_train_function([force])

Creates a function that executes one step of training.

predict(x[, batch_size, verbose, steps, ...])

Generates output predictions for the input samples.

predict_generator(generator[, steps, ...])

Generates predictions for the input samples from a data generator.

predict_on_batch(x)

Returns predictions for a single batch of samples.

predict_step(data)

The logic for one inference step.

reset_metrics()

Resets the state of all the metrics in the model.

reset_states()

save(filepath[, overwrite, save_format])

Saves a model as a TensorFlow SavedModel or HDF5 file.

save_own_variables(store)

Saves the state of the layer.

save_spec([dynamic_batch])

Returns the tf.TensorSpec of call args as a tuple (args, kwargs).

save_weights(filepath[, overwrite, ...])

Saves all layer weights.

set_params(**params)

Set the parameters of this learner.

set_weights(weights)

Sets the weights of the layer, from NumPy arrays.

summary([line_length, positions, print_fn, ...])

Prints a string summary of the network.

test_on_batch(x[, y, sample_weight, ...])

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.

to_yaml(**kwargs)

Returns a yaml 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.

with_name_scope(method)

Decorator to automatically enter the module name scope.

Attributes

activity_regularizer

Optional regularizer function for the output of this layer.

autotune_steps_per_execution

Settable property to enable tuning for steps_per_execution

compute_dtype

The dtype of the layer's computations.

distribute_reduction_method

The method employed to reduce per-replica values during training.

distribute_strategy

The tf.distribute.Strategy this model was created under.

dtype

The dtype of the layer weights.

dtype_policy

The dtype policy associated with this layer.

dynamic

Whether the layer is dynamic (eager-only); set in the constructor.

inbound_nodes

Return Functional API nodes upstream of this layer.

input

Retrieves the input tensor(s) of a layer.

input_mask

Retrieves the input mask tensor(s) of a layer.

input_shape

Retrieves the input shape(s) of a layer.

input_spec

InputSpec instance(s) describing the input format for this layer.

jit_compile

Specify whether to compile the model with XLA.

layers

losses

List of losses added using the add_loss() API.

metrics

Return metrics added using compile() or add_metric().

metrics_names

Returns the model's display labels for all outputs.

my_params

name

Name of the layer (string), set in the constructor.

name_scope

Returns a tf.name_scope instance for this class.

non_trainable_variables

Sequence of non-trainable variables owned by this module and its submodules.

non_trainable_weights

List of all non-trainable weights tracked by this layer.

outbound_nodes

Return Functional API nodes downstream of this layer.

output

Retrieves the output tensor(s) of a layer.

output_mask

Retrieves the output mask tensor(s) of a layer.

output_shape

Retrieves the output shape(s) of a layer.

run_eagerly

Settable attribute indicating whether the model should run eagerly.

state_updates

Deprecated, do NOT use!

stateful

steps_per_execution

Settable `steps_per_execution variable. Requires a compiled model.

submodules

Sequence of all sub-modules.

supports_masking

Whether this layer supports computing a mask using compute_mask.

trainable

trainable_variables

Sequence of trainable variables owned by this module and its submodules.

trainable_weights

List of all trainable weights tracked by this layer.

updates

variable_dtype

Alias of Layer.dtype, the dtype of the weights.

variables

Returns the list of all layer variables/weights.

weights

Returns the list of all layer variables/weights.
__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)[source]
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)[source]

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)[source]

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.

  1. Assembling the PDE residual using these derivatives.

  2. Calculating the mean squared error of the residual.

  3. 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)[source]

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()[source]

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)[source]

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' )