models.pinn package

Submodules

models.pinn.boundary_condition module

class models.pinn.boundary_condition.BoundaryCondition(points)[source]

Bases: ABC

Abstract base class for defining boundary conditions.

Parameters:

points (Tensor) – The points where the boundary condition is defined.

_points

The points where the boundary condition is defined.

Type:

Tensor

abstract loss(pred)[source]

Computes the loss for the given prediction.

Parameters:

pred (Tensor) – The predicted values.

Returns:

The loss value.

Return type:

Tensor

property points

models.pinn.pinn module

class models.pinn.pinn.OldPINN(neural_net, device)[source]

Bases: ABC

This class represents a Physics-Informed Neural Network (PINN) model.

Parameters:

  • neural_net (torch.nn.Module) – The neural network model.

  • device (str) – The device to run the model on (e.g., ‘cpu’, ‘cuda’).

device

The device the model is running on.

Type:

str

model

The neural network model.

Type:

torch.nn.Module

bc_data_loss(pred, y, boundary_conditions, use_bfloat16=False)[source]

Computes the loss from boundary conditions and data.

Parameters:

  • pred (torch.Tensor) – The predicted output tensor.

  • y (torch.Tensor) – The target tensor.

  • boundary_conditions (List[BoundaryCondition]) – The list of boundary conditions.

  • use_bfloat16 (bool, optional) – Whether to use bfloat16 precision. Defaults to False.

Returns:

The list of loss tensors.

Return type:

List[torch.Tensor]

compute_loss(x, y, boundary_conditions, use_bfloat16=False)[source]

Computes the total loss for training.

Parameters:

  • x (torch.Tensor) – The input tensor.

  • y (torch.Tensor) – The target tensor.

  • boundary_conditions (List[BoundaryCondition]) – The list of boundary conditions.

  • use_bfloat16 (bool, optional) – Whether to use bfloat16 precision. Defaults to False.

Returns:

The list of loss tensors.

Return type:

List[torch.Tensor]

abstract pde_loss(pred, *input_variables)[source]

Computes the loss from the partial differential equation (PDE).

Parameters:

  • pred (torch.Tensor) – The predicted output tensor.

  • *input_variables (torch.Tensor) – The input variables for the PDE. e.g. x, y, t.

Returns:

The loss tensor.

Return type:

torch.Tensor

plot_training_logs(logs)[source]

Plots the training logs.

Parameters:

logs (dict) – The training logs.

train_model(optimizer, epochs, x_train, y_train=None, boundary_conditions=[], print_every=100, loaded_logs=None, batch_size=None, x_test=None, y_test=None, use_bfloat16=False)[source]

Trains the PINN model.

Parameters:

  • optimizer (torch.optim.Optimizer) – The optimizer used for the training.

  • epochs (int) – The number of epochs to train for.

  • x_train (torch.Tensor) – The input training tensor containing the collocation points.

  • y_train (torch.Tensor, optional) – The target training tensor if any simulation data wants to be given. Defaults to None.

  • boundary_conditions (List[BoundaryCondition], optional) – The list of boundary conditions. Defaults to [].

  • print_every (int, optional) – The interval for printing progress. Defaults to 100.

  • loaded_logs (dict, optional) – Loaded training logs to be used as initial logs. Defaults to None.

  • batch_size (int, optional) – The batch size. If none, the batch size will be equal to the number of collocation points given on x_train Defaults to None.

  • x_test (torch.Tensor, optional) – The input test tensor containg the test points. Defaults to None.

  • y_test (torch.Tensor, optional) – The target test tensor with simulation data on the test points. If no value is given, then the test loss will be the loss o the PDE + the loss of the boundary conditions. Defaults to None.

  • use_bfloat16 (bool, optional) – Whether to use bfloat16 precision. Defaults to False.

Returns:

The training logs.

Return type:

dict

class models.pinn.pinn.PINN(neural_net, device)[source]

Bases: ABC

This class represents a Physics-Informed Neural Network (PINN) model.

Parameters:

  • neural_net (torch.nn.Module) – The neural network model.

  • device (str) – The device to run the model on (e.g., ‘cpu’, ‘cuda’).

device

The device the model is running on.

Type:

str

model

The neural network model.

Type:

torch.nn.Module

bc_data_loss(pred, y, boundary_conditions, use_bfloat16=False)[source]

Computes the loss from boundary conditions and data.

Parameters:

  • pred (torch.Tensor) – The predicted output tensor.

  • y (torch.Tensor) – The target tensor.

  • boundary_conditions (List[BoundaryCondition]) – The list of boundary conditions.

  • use_bfloat16 (bool, optional) – Whether to use bfloat16 precision. Defaults to False.

Returns:

The list of loss tensors.

Return type:

List[torch.Tensor]

compute_loss(x, y, boundary_conditions, use_bfloat16=False)[source]

Computes the total loss for training.

Parameters:

  • x (torch.Tensor) – The input tensor.

  • y (torch.Tensor) – The target tensor.

  • boundary_conditions (List[BoundaryCondition]) – The list of boundary conditions.

  • use_bfloat16 (bool, optional) – Whether to use bfloat16 precision. Defaults to False.

Returns:

The list of loss tensors.

Return type:

List[torch.Tensor]

fit(train_dataset, optimizer_class=<class 'torch.optim.adam.Adam'>, optimizer_params={}, lr_scheduler_class=None, lr_scheduler_params={}, epochs=1000, boundary_conditions=[], update_logs_steps=1, loaded_logs=None, batch_size=None, eval_dataset=None, use_bfloat16=False, **kwargs)[source]

Trains the PINN model.

Parameters:

  • train_dataset (BaseDataset) – The training dataset.

  • optimizer_class (torch.optim.Optimizer, optional) – The optimizer class. Defaults to torch.optim.Adam.

  • optimizer_params (dict, optional) – The optimizer parameters. Defaults to {}.

  • lr_scheduler_class (torch.optim.lr_scheduler._LRScheduler, optional) – The learning rate scheduler class. Defaults to None.

  • lr_scheduler_params (dict, optional) – The learning rate scheduler parameters. Defaults to {}.

  • epochs (int, optional) – The number of epochs to train for. Defaults to 1000.

  • boundary_conditions (List[BoundaryCondition], optional) – The list of boundary conditions. Defaults to [].

  • update_logs_steps (int, optional) – The interval for updating the progress. Defaults to 100.

  • loaded_logs (dict, optional) – Loaded training logs to be used as initial logs. Defaults to None.

  • batch_size (int, optional) – The batch size. If none, the batch size will be equal to the number of collocation points given on train_dataset. Defaults to None.

  • eval_dataset (BaseDataset, optional) – The evaluation dataset. Defaults to None.

  • use_bfloat16 (bool, optional) – Whether to use bfloat16 precision. Defaults to False.

  • **kwargs – Additional keyword arguments.

Returns:

The training logs.

Return type:

dict

classmethod load(path, device='cpu')[source]

Loads the model from a file.

Parameters:

  • path (str) – The path to load the model.

  • neural_net (torch.nn.Module) – The neural network model.

  • device (str, optional) – The device to run the model on. Defaults to ‘cpu’.

Returns:

The loaded PINN model.

Return type:

PINN

abstract pde_loss(pred, *input_variables)[source]

Computes the loss from the partial differential equation (PDE).

Parameters:

  • pred (torch.Tensor) – The predicted output tensor.

  • *input_variables (torch.Tensor) – The input variables for the PDE. e.g. x, y, t.

Returns:

The loss tensor.

Return type:

torch.Tensor

plot_training_logs(logs)[source]

Plots the training logs.

Parameters:

logs (dict) – The training logs.

predict(X, **kwargs)[source]

Predicts for the input dataset.

Parameters:

X (BaseDataset) – The input dataset.

Returns:

The predictions of the model.

Return type:

np.ndarray

save(path)[source]

Saves the model to a file using torchscript.

Parameters:

path (str) – The path to save the model.

models.pinn.piratenet module

class models.pinn.piratenet.PirateNet(input_dim, output_dim, num_blocks, hidden_dim=256, s=1.0, activation=<function tanh>)[source]

Bases: Module

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

initialize_last_layer(Y, input_data)[source]
initialize_weights()[source]
class models.pinn.piratenet.PirateNetBlock(hidden_dim)[source]

Bases: Module

forward(x, u, v)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

models.pinn.simple_mlp module

class models.pinn.simple_mlp.Net(layer_dim, activation=<function tanh>, device=device(type='cpu'), **kwargs)[source]

Bases: Module, Model

Initialize the MLP model.

Parameters:

  • layer_dim (list) – List containing the number of neurons in each layer.

  • activation (function) – Activation function to use. Default is F.tanh.

  • device (torch.device) – Device to use for computation. Default is torch.device(“cpu”).

  • **kwargs – Additional keyword arguments.

classmethod create_optimized_model(train_dataset, eval_dataset, optuna_optimizer)[source]

Create an optimized model using Optuna.

Parameters:

  • train_dataset (BaseDataset) – The training dataset.

  • eval_dataset (Optional[BaseDataset]) – The evaluation dataset.

  • optuna_optimizer (OptunaOptimizer) – The optimizer to use for optimization.

Returns:

A tuple containing the optimized model and the best parameters for training found by the optimizer

Return type:

Tuple[Model, Dict]

fit(train_dataset, eval_dataset=None, epochs=100, batch_size=32, lr=0.001, loss_fn=MSELoss(), optimizer_class=<class 'torch.optim.adam.Adam'>, print_rate=1, **kwargs)[source]

Fit the model to the training data.

Parameters:

  • train_dataset (BaseDataset) – The training dataset.

  • eval_dataset (Optional[BaseDataset]) – The evaluation dataset. Default is None.

  • epochs (int) – The number of epochs to train. Default is 100.

  • batch_size (int) – The batch size. Default is 32.

  • lr (float) – The learning rate. Default is 0.001.

  • loss_fn (nn.Module) – The loss function. Default is nn.MSELoss().

  • optimizer_class (torch.optim.Optimizer) – The optimizer class. Default is torch.optim.Adam.

  • print_rate (int) – The rate at which to print the training loss. Default is 1.

  • **kwargs – Additional keyword arguments.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

static load(path)[source]

Load the model from a file.

Parameters:

path (str) – The path to load the model from.

Returns:

The loaded model.

Return type:

Model

Examples

>>> model = Net.load("model.pth")
>>> model.predict(X)
predict(X, rescale_output=True)[source]

Predict the target values for the input data.

Parameters:

  • X (BaseDataset) – The input data.

  • rescale_output (bool) – Whether to rescale the output with the scaler of the dataset. Default is True.

Returns:

The predicted target values.

Return type:

torch.Tensor

save(path)[source]

Save the model to a file.

Parameters:

path (str) – The path to save the model.

Module contents