CoupledNODE

cnn(; setup, radii, channels, activations, use_bias, channel_augmenter = identity, rng = Random.default_rng())

Constructs a convolutional neural network model closure(u, θ) that predicts the commutator error (i.e. closure).

Arguments

• setup: IncompressibleNavierStokes.jl setup
• radii: An array (size n_layers) with the radii of the kernels for the convolutional layers. Kernels will be symmetrical of size 2r+1.
• channels: An array (size n_layers) with channel sizes for the convolutional layers.
• activations: An array (size n_layers) with activation functions for the convolutional layers.
• use_bias: An array (size n_layers) with booleans indicating whether to use bias in each convolutional layer.
• rng: A random number generator (default: Random.default_rng()).

Returns

A tuple (chain, params, state) where

• chain: The constructed Lux.Chain model.
• params: The parameters of the model.
• state: The state of the model.

create_callback(model, val_io_data; lhist=[], lhist_train=[], nunroll=10, rng=rng, plot_train=true)

Create a callback function for training and validation of a model.

Arguments

• model: The model for the rhs.
• θ: parameters of the model (trainable).
• val_io_data: The validation io_array.
• loss_function: The loss function to be used.
• st: The state of the model.
• callbackstate: a NamedTuple that is updated durign the trainign and contains:
  • θmin: The parameters at the minimum validation loss.
  • loss_min: The minimum validation loss.
  • lhist_val: A list to store the history of validation losses. Defaults to a new empty list.
  • lhist_train: A list to store the history of training losses. Defaults to a new empty list.
• nunroll: The number of unroll steps for the validation loss. It does not have to be the same as the loss function! Pertinent for a-posteriori training.
• rng: The random number generator to be used.
• plot_train: A boolean flag to indicate whether to plot the training loss.
• do_plot: A boolean flag to indicate whether to generate the plots. In HPC systems we may want to deactivate it.
• plot_every: The frequency of plotting the loss history. Defaults to 10. The loss is also averaged in this window.

Returns

A NamedTuple with the callbackstate` and the callback function. The callback function is used during training to compute and log validation and training losses, and optionally plot the loss history.

Callback function arguments

• p: The parameters of the model at the current training step.
• l_train: The training loss at the current training step.

Callback function returns

• callbackstate: the updated instance of the callback state.

create_loss_post_lux(rhs; sciml_solver = Tsit5(), kwargs...)

Creates a loss function for a-posteriori fitting using the given right-hand side (RHS) function rhs. The loss function computes the sum of squared differences between the predicted values and the actual data, normalized by the sum of squared actual data values.

Arguments

• rhs::Function: The right-hand side function for the ODE problem.
• sciml_solver::AbstractODEAlgorithm: (Optional) The SciML solver to use for solving the ODE problem. Defaults to Tsit5().
• kwargs...: Additional keyword arguments to pass to the solver.

Returns

• loss_function::Function: A function that computes the loss given a model, parameters, state, and data (u, t).

loss_function Arguments

• model: The model to evaluate.
• ps: Parameters for the model.
• st: State of the model.
• (u, t): Tuple containing the data u and time points t.

loss_function Returns

• loss: The computed loss value.
• st: The model state (unchanged).
• metadata::NamedTuple: A named tuple containing the predicted values y_pred.

This makes it compatible with the Lux ecosystem.

Wrap loss function loss(batch, θ).

The function loss should take inputs like loss(f, x, y, θ).

[DEPRECATED] createrandlossMulDtO(target; nunroll, nintervals=1, nsamples, λc, λl1)

Creates a random loss function for the multishooting method with multiple shooting intervals.

Arguments

• target: The target data for the loss function.
• training_CNODE: Model CNODE.
• st: state of the neural part.
• nunroll: The number of time steps to unroll.
• noverlaps: The number of time steps that overlaps between consecutive intervals.
• nintervals: The number of shooting intervals.
• nsamples: The number of samples to select.
• λ_c: The weight for the continuity term. It sets how strongly we make the pieces match (continuity term).
• λ_l1: The coefficient for the L1 regularization term in the loss function.

Returns

• randloss_MulDtO: A random loss function for the multishooting method.

create_randloss_derivative(input_data, F_target, f, st; n_use = size(input_data, 2), λ=0, λ_c = 0)

Create a randomized loss function that compares the derivatives. This is done because using the entire dataset at each iteration would be too expensive.

Arguments

• input_data: The input data.
• F_target: The target data.
• f: The model function.
• st: The model state.
• n_use: The number of samples to use for the loss function. Defaults to the number of samples in the input data.
• λ and λ_c: The regularization parameter. Defaults to 0.

Returns

A function that computes the mean squared error loss and takes as input the model parameters.

Interpolate closure force from volume centers to volume faces.

Interpolate velocity components to volume centers.

TODO, D and dir can be parameters istead of arguments I think

[DEPRECATED] lossMulDtOoneset(trajectory, θ; λc=1e1, λl1=1e1, nunroll, nintervals, nsamples=nsamples)

Compute the loss function for the multiple shooting method with a continuous neural ODE (CNODE) model. Check https://docs.sciml.ai/DiffEqFlux/dev/examples/multiple_shooting/ for more details.

Arguments

• trajectory: The trajectory of the system.
• θ: The parameters of the CNODE model.
• training_CNODE: Model CNODE.
• λ_c: The weight for the continuity term. It sets how strongly we make the pieces match (continuity term). Default is 1e1.
• λ_l1: The weight for the L1 regularization term. Default is 1e1.
• nunroll: The number of time steps to unroll the trajectory.
• noverlaps: The number of time steps that overlaps between each consecutive intervals.
• nintervals: The number of intervals to divide the trajectory into.
• nsamples: The number of samples.

Returns

• loss: The computed loss value.

loss_posteriori_optim(model, p, st, dataloader)

Compute the loss function for a-posteriori fitting using the given model, parameters, state, and dataloader. This function is compatible with training via Optim.jl.

Arguments

• model: The model to evaluate.
• p: Parameters for the model.
• st: State of the model.
• dataloader::Function: A function that returns the data (u, t).

Returns

• loss: The computed loss value.

loss_priori_lux(model, ps, st, (x, y))

Compute the loss as the sum of squared differences between the predicted values and the target values, normalized by the sum of squared target values. Format of function signature and outputs are compatible with the Lux ecosystem.

Arguments

• model: The model to be evaluated.
• ps: Parameters for the model.
• st: State of the model.
• (x, y): A tuple where x is the input data and y is the target data.

Returns

• loss: The computed loss value.
• st_: The updated state of the model.
• (; y_pred = y_pred): Named tuple containing the predicted values y_pred.

Compute MSE between f(x, θ, st) and y. The MSE is further divided by normalize(y). This signature has been chosen for compatibility with Lux.

mean_squared_error(f, st, x, y, θ, λ)

Random a priori loss function. Use this function to train the closure term to reproduce the right hand side.

Arguments:

• f: Function that represents the model.
• st: State of the model.
• x: Input data.
• y: Target data.
• θ: Parameters of the model.
• λ and λ_c: Regularization parameters.

Returns:

• total_loss: Mean squared error loss.

validate_results(model, p, dataloader, nuse = 100)

Validate the results of the model using the given parameters p, dataloader, and number of samples nuse.

Arguments

• model: The model to evaluate.
• p: Parameters for the model.
• dataloader::Function: A function that returns the data (u, t).
• nuse::Int: The number of samples to use for validation. Defaults to 100.

Returns

• loss: The computed loss value.

NN_padded_to_NN_nopad(u, setup)

Creates a view of the input velocity field u from the neural network data style u[n, n, D, batch] but without boundaries.

Arguments

• u: Velocity field in NN style.
• setup: IncompressibleNavierStokes.jl setup.

Returns

• u: Velocity field view without boundaries.

create_dataloader_posteriori(io_array; nunroll=10, device=identity, rng)

Creates a dataloader function for a-posteriori fitting from the given io_array.

Arguments

• io_array: A structure containing the data arrays u and t.
• The io_array.u array is expected to have dimensions (nx, ny... , dim, samples, nt).
• nunroll: The number of time steps to unroll (default is 10).
• rng: A random number generator.

Returns

• A function dataloader that, when called, returns a tuple with:
  • u: A view into the u array of io_array for a randomly selected sample and time steps.
  • t: The corresponding time steps from the t array of io_array.

Notes

• The nt dimension must be greater than or equal to nunroll.
• Have only tested in 2D cases.
• It assumes that the data are loaded in batches of size 1

create_dataloader_prior(io_array; nunroll=10, device=identity, rng)

Create dataloader that uses a batch of `batchsize` random samples from

data at each evaluation.

Arguments

• io_array: An named tuple with the data arrays u and c.
• batchsize: The number of samples to use in each batch.
• device: A function to move the data to the desired device (e.g. gpu).

Returns

• A function dataloader that, when called, returns a tuple with:
  • u: baru from `ioarray` (input to NN)
  • c: commutator error from io_array (label)

create_io_arrays_posteriori(data, setups)

Main differences between this function and NeuralClosure.createioarrays

• we do not need the commutator error c.
• we need the time and we keep the initial condition.
• put time dimension in the end, since SciML also does.

Returns

A named tuple with fields u and t. u is a matrix without padding and shape (nless..., D, sample, t)

create_io_arrays_priori(data, setups)

Create $(\bar{u}, c)$ pairs for training.

Returns

A named tuple with fields u and c. (without boundary conditions padding)

create_right_hand_side(setup, psolver)

Create right hand side function f(u, p, t) compatible with the OrdinaryDiffEq ODE solvers. Note that u has to be an array. To convert the tuple u = (ux, uy) to an array, use stack(u).

create_right_hand_side_inplace(setup, psolver)

In place version of create_right_hand_side.

create_right_hand_side_with_closure(setup, psolver, closure, st)

Create right hand side function f(u, p, t) compatible with SciML ODEProblem. This formulation was the one proved best in Syver's paper i.e. DCF. u has to be an array in the NN style e.g. [n , n, D], with boundary conditions padding.