neural_operators.nn#

neural_operators.nn.nn_util.get_surrogate_specs(name, data_prefix, model_path)[source]#
neural_operators.nn.nn_util.load_data_and_model(name, data_prefix, model_path)[source]#
neural_operators.nn.nn_util.load_surrogate_model(name, data_prefix, model_path, model, u_comps=None)[source]#
class neural_operators.nn.mlp.torch_mlp.MLP(input_size, hidden_size, num_classes, depth, act)[source]#

Bases: Module

Plain fully-connected net, depth layers total. Note: forward(x, final_act=True) applies relu (not act) to the output layer instead of leaving it linear.

forward(x, final_act=False)[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.

class neural_operators.nn.deeponet.torch_deeponet.DeepONet(num_layers, num_neurons, act, num_br_outputs, num_tr_outputs, num_inp_fn_points, out_coordinate_dimension, num_Y_components, save_file=None)[source]#

Bases: Module

Branch net encodes the input function, trunk net encodes the output coordinates; forward combines them as a dot product plus bias. num_Y_components > 1 splits the branch output into that many blocks, one dot product each, for vector-valued output.

convert_np_to_tensor(array)[source]#
forward(X, X_trunk)[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.

predict(X, X_trunk)[source]#
train(train_data, test_data, batch_size=32, epochs=1000, lr=0.001, log=True, loss_print_freq=100, save_model=False, save_file=None, save_epoch=100)[source]#

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

neural_operators.nn.deeponet.load_data_and_deeponet.load_data_and_deeponet(data_path, model_path)[source]#
class neural_operators.nn.pcanet.torch_pcanet.PCANet(num_layers, num_neurons, act, num_inp_red_dim, num_out_red_dim, save_file=None)[source]#

Bases: Module

Plain MLP mapping PCA-reduced input coefficients to PCA-reduced output coefficients; the PCA encode/decode itself lives in DataProcessor, not here.

convert_np_to_tensor(array)[source]#
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.

predict(X)[source]#
train(train_data, test_data, batch_size=32, epochs=1000, lr=0.001, log=True, loss_print_freq=100, save_model=False, save_file=None, save_epoch=100)[source]#

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

neural_operators.nn.pcanet.load_data_and_pcanet.load_data_and_pcanet(data_path, model_path)[source]#
class neural_operators.nn.fno.torch_fno2d.FNO2D(num_layers, width, fourier_modes1, fourier_modes2, num_Y_components, save_file=None)[source]#

Bases: Module

Input-projects (m(x,y), x, y) to width channels, applies num_layers stacked FNO2DLayers, output-projects to num_Y_components. Operates on a fixed regular grid, not a mesh, unlike DeepONet/PCANet.

convert_np_to_tensor(array)[source]#
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.

predict(X)[source]#
train(train_data, test_data, batch_size=32, epochs=1000, lr=0.001, log=True, loss_print_freq=100, save_model=False, save_file=None, save_epoch=100)[source]#

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Parameters:

mode (bool) – whether to set training mode (True) or evaluation mode (False). Default: True.

Returns:

self

Return type:

Module

class neural_operators.nn.fno.torch_fno2dlayer.FNO2DLayer(in_channels, out_channels, modes1, modes2, apply_act=True, act=<built-in function gelu>)[source]#

Bases: Module

One Fourier layer: spectral conv (keeps the lowest modes1 x modes2 Fourier modes) plus a pointwise 1x1 conv skip connection, then an activation.

compl_mul2d(a, b)[source]#
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.

fourier_transform(x)[source]#
linear_transform(x)[source]#
neural_operators.nn.fno.load_data_and_fno.load_data_and_fno(data_path, model_path)[source]#