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_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

load_checkpoint(checkfile)

Load a training checkpoint from the specified file.

Arguments

• checkfile::String: The path to the checkpoint file.

Returns

• callbackstate: The state of the callback at the checkpoint.
• trainstate: The state of the training process at the checkpoint.
• epochs_trained::Int: The number of epochs completed at the checkpoint.

Example

callbackstate, trainstate, epochs_trained = load_checkpoint("checkpoint.jld2")

[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.