Reference#

IRK_PINNs#

class IRK_PINNs.runge_kutta.MLP(feats_num, lower_bound, upper_bound, activ_func=<function bipolar_sigmoid>, parent=<flax.linen.module._Sentinel object>, name=None)#

MLP to be used for ‘t’ and ‘s’ functions in RNVP model.

Parameters:

feats_num (Sequence[int]) – List of integers representing the number of neurons in each layer.
lower_bound (jnp.ndarray) – Lower bound of the input data.
upper_bound (jnp.ndarray) – Upper bound of the input data.
activ_func (Callable, optional) – Activation function to be used in the neural network. Defaults to nn.tanh.
parent (Type[Module] | Type[Scope] | Type[_Sentinel] | None) –
name (str | None) –

class IRK_PINNs.runge_kutta.PINN_RK(hidden_layers, dt, q, n, lower_bound, upper_bound, diff_eq, activ_func=<function bipolar_sigmoid>, kwargs=None, parent=<flax.linen.module._Sentinel object>, name=None)#

Follows a discrete time inference according to the paradigm outlined in https://arxiv.org/pdf/1711.10561.pdf. We proceed to use their conventions.

Parameters:

hidden_layers (List[int]) – List of integers, number of neurons in each layer.
dt (float) – Corresponds to $Delta t$ in the paper, represents time difference between where the data is sampled and where the model predicts.
q (int) – Corresponds to $q$ in the paper, number of ‘steps’ to be used in the RK method.
n (int) – Not in the paper, number of output dimensions (in our case, the number of coordinates of the phase space).
lower_bound (float) – Lower bound of the input data.
upper_bound (float) – Upper bound of the input data.
diff_eq (Callable) – Differential equation to be used in the model. Input should have shape the neural network used for the priors, the input data to evaluate, q, and n (need to pass them to separate arrays) and output should have shape (n, N, q). For more information, see create_AC_diff_eq function above.
activ_func (Callable) – Activation function to be used in the neural network.
kwargs (dict | None) – Additional arguments to be passed to the differential equation.
parent (Type[Module] | Type[Scope] | Type[_Sentinel] | None) –
name (str | None) –

setup()#

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

Immediately when invoking apply(), init() or init_and_output().
Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):
class MyModule(nn.Module):
  def setup(self):
    submodule = Conv(...)

    # Accessing `submodule` attributes does not yet work here.

    # The following line invokes `self.__setattr__`, which gives
    # `submodule` the name "conv1".
    self.conv1 = submodule

    # Accessing `submodule` attributes or methods is now safe and
    # either causes setup() to be called once.
Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

Return type:: None

IRK_PINNs.runge_kutta.create_AC_diff_eq(neural_net, X, q, n)#

Example of implementation of a differential equation and its attributes. Creates the array N for the differential equation of the AC model.

In the following, the conventions used are: - n is the dimension of the output. - q is the number of steps in the RK method. - m is the dimension of the input. - N is the number of points in a batch.

Parameters:

neural_net (Callable) – Neural Network to be used in the model.
X (jnp.ndarray) – Array of shape (N, m) of the input data (x^{n, i} in the paper).
q (int) – Number of ‘steps’ to be used in the RK method.
n (int) – Number of output dimensions (in our case, the number of coordinates of the phase space).

Returns:

Array of shape (n, N, q) of the corresponding differential equations (mathcal{N}[u^{n+c_i}] in the paper).

Return type:

jnp.ndarray

Note

The dimensions of the variables to be created are as follows: - U1_sep_ (jnp.ndarray): This array has shape (n, N, q) and contains the values of the output of the neural network, which are u^{n+c_i} in the paper. - dU1_sep_ (jnp.ndarray): This array has shape (n, N, q, m) and contains the derivatives with respect to each coordinate of the u^{n+c_i}. - ddU1_sep_ (jnp.ndarray): This array has shape (n, N, q, m, m) and contains second derivatives.

IRK_PINNs.runge_kutta.dir_ = '/builds/CMI/CMI-public/runge-kutta-pinn/IRK_PINNs'#

IRK-PINNs.runge_kutta.py

This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.

IRK_PINNs.runge_kutta.loss_RK(params, model, x, u_training, q, loss_, boundary_conditions)#

Loss function for PINN_RK model using conventions from https://arxiv.org/pdf/1711.10561.pdf

Parameters:

params – parameters of the model
model – PINN_RK model
x – input data (x^{n, i} in the paper)
u_training – output data (array of output data if output has more than one dimension) (u^{n, i} in the paper)

Returns:

loss value

Return type:

loss

IRK_PINNs.runge_kutta.train_ADAM(model, x, u_training, q, epochs, lr, boundary_conditions=None, test_epoch=5000, evaluation_data=None)#

Train the PINN_RK model using second-order LBFGS optimizer

Parameters:

model (nn.Module) – PINN_RK model.
params (FrozenDict) – Frozen dictionary with parameters, obtained from Adam training.
x (jnp.array) – Input data (x^{n, i} in the paper).
u_training (jnp.array) – Output data (array of output data if the output has more than one dimension) (u^{n, i} in the paper). Should have shape (N, n).
q (int) – Corresponds to $q$ in the paper, the number of ‘steps’ to be used in the RK method.
epochs (int) – Number of epochs to train the model.
boundary_conditions (Callable[[jnp.ndarray, jnp.ndarray], float], optional) – Function that returns the value of the boundary conditions.
test_epoch (int, optional) – The number of epochs in which to perform a print of the training status.
evaluation_data (Tuple, optional) – Set of phase space points and solutions for them to be used as L^2 approximation.
tol (float, optional) – Tolerance of the second-order optimizer to be stopped.
lr (float) –

Returns:

Trained parameters of the model.

IRK_PINNs.runge_kutta.train_LBFGS(model, params, x, u_training, q, epochs, boundary_conditions=None, test_epoch=5000, evaluation_data=None, tol=1e-18)#

Train the PINN_RK model using second order LBFGS optimizer

Parameters:

model (Module) – PINN_RK model
params (FrozenDict) – Frozen dictionary with parameters, obtained from Adam training
x (array) – input data (x^{n, i} in the paper)
u_training (array) – output data (array of output data if output has more than one dimension) (u^{n, i} in the paper) should have shape (N, n)
q (int) – corresponds to $q$ in the paper, number of ‘steps’ to be used in the RK method
epochs (int) – number of epochs to train the model
boundary_conditions (Callable[[Array, Array], float] | None) – function that returns the value of the boundary conditions
test_epoch (int) – the number of epochs in which we want to perform a print of the status of the training
evaluation_data (Tuple | None) – set of phase space points and solutions for them to be used as L^2 approximation
tol (float) – tolerance of the second order optimizer to be stopped

Returns: trained parameters of the model

IRK_PINNs.runge_kutta.train_RK(model, x, u_training, q, epochs, lr=0.0001, boundary_conditions=None, test_epoch=5000, evaluation_data=None, epochs_LBFGS=None, test_epoch_LBFGS=None, tol=1e-12)#

Train a PINN using the discrete time scheme (Runge Kutta of qth order) using optimization with ADAM followed by LBFGS.

Parameters:

model (nn.Module) – PINN model specified above.
x (jnp.array) – Input data (x^{n, i} in the paper).
u_training (jnp.array) – Output data (array of output data if output has more than one dimension) (u^{n, i} in the paper). Should have shape (N, n).
q (int) – Corresponds to $q$ in the paper, the number of ‘steps’ to be used in the RK method.
epochs (int) – Number of epochs to train the model with ADAM.
lr (float, optional) – Learning rate for the ADAM optimizer. Defaults to 1e-4.
boundary_conditions (Callable[[jnp.ndarray, jnp.ndarray], float], optional) – Function that returns the value of the boundary conditions. Defaults to None.
test_epoch (int, optional) – Loss and error are printed out on multiples of this variable. Defaults to 5000.
evaluation_data (Tuple, optional) – Exact data for the model to evaluate against for the real error. Defaults to None (no real error is printed).
epochs_LBFGS (int, optional) – Number of epochs for LBFGS optimizer. Defaults to None (uses the same as for ADAM).
test_epoch_LBFGS (int, optional) – Test epoch for LBFGS. Defaults to None (uses the same as for ADAM).
tol (float, optional) – Tolerance for LBFGS optimizer. Defaults to 1e-12.

Returns:

Trained parameters.

Return type:

FrozenDict