PhysicsInformedNN Discretizer for PDESystems

PhysicsInformedNN(
    chain, strategy; init_params = nothing, init_states = nothing,
    phi = nothing, param_estim = false, additional_loss = nothing,
    adaptive_loss = nothing, logger = nothing, log_options = LogOptions(),
    iteration = nothing, kwargs...
)

A discretize algorithm for the ModelingToolkit PDESystem interface, which transforms a PDESystem into an OptimizationProblem using the Physics-Informed Neural Networks (PINN) methodology.

Positional Arguments

• chain: a vector of Lux/Flux chains with a d-dimensional input and a 1-dimensional output corresponding to each of the dependent variables. Note that this specification respects the order of the dependent variables as specified in the PDESystem. Flux chains will be converted to Lux internally using adapt(FromFluxAdaptor(), chain).
• strategy: determines which training strategy will be used. See the Training Strategy documentation for more details.

Keyword Arguments

• init_params: the initial parameters of the neural networks. If init_params is not given, then the neural network default parameters are used. Note that for Lux, the default will convert to Float64.
• init_states: the initial states of the neural networks. If init_states is not given, then the neural network default states are used. Note that for Lux, the default will convert to Float64.
• phi: a trial solution, specified as phi(x,p) where x is the coordinates vector for the dependent variable and p are the weights of the phi function (generally the weights of the neural network defining phi). By default, this is generated from the chain. This should only be used to more directly impose functional information in the training problem, for example imposing the boundary condition by the test function formulation.
• adaptive_loss: the choice for the adaptive loss function. See the adaptive loss page for more details. Defaults to no adaptivity.
• additional_loss: a function additional_loss(phi, θ, p_) where phi are the neural network trial solutions, θ are the weights of the neural network(s), and p_ are the hyperparameters of the OptimizationProblem. If param_estim = true, then θ additionally contains the parameters of the differential equation appended to the end of the vector.
• param_estim: whether the parameters of the differential equation should be included in the values sent to the additional_loss function. Defaults to false.
• logger: a logging object (e.g. a TensorBoardLogger) used for recording loss values and adaptive weights during training. Defaults to nothing (no logging).
• log_options: a LogOptions struct controlling logging frequency (e.g. how often to write loss values to the logger). Separate from logger to allow configuring log frequency independently of the logger type.
• iteration: an optional external iteration counter (a Ref{Int} or Vector{Int} of length 1) shared with the caller so the caller can read or control the training step count. If not provided, an internal counter is created and auto-incremented.
• kwargs: Extra keyword arguments which are splatted to the OptimizationProblem on solve.

An encoding of the test function phi that is used for calculating the PDE value at domain points x

Fields:

• f: A representation of the chain function.
• st: The state of the Lux.AbstractLuxLayer. It should be updated on each call.

prob = discretize(pde_system::PDESystem, discretization::PhysicsInformedNN)

Transforms a symbolic description of a ModelingToolkit-defined PDESystem and generates an OptimizationProblem for Optimization.jl whose solution is the solution to the PDE.

prob = symbolic_discretize(pde_system::PDESystem, discretization::AbstractPINN)

symbolic_discretize is the lower level interface to discretize for inspecting internals. It transforms a symbolic description of a ModelingToolkit-defined PDESystem into a PINNRepresentation which holds the pieces required to build an OptimizationProblem for Optimization.jl or a Likelihood Function used for HMC based Posterior Sampling Algorithms AdvancedHMC.jl which is later optimized upon to give Solution or the Solution Distribution of the PDE.

For more information, see discretize and PINNRepresentation.

PINNRepresentation`

An internal representation of a physics-informed neural network (PINN). This is the struct used internally and returned for introspection by symbolic_discretize.

Fields

• eqs: The equations of the PDE

• bcs: The boundary condition equations

• domains: The domains for each of the independent variables

• eq_params: The symbolic parameters of the PDE system (e.g. physical constants to be estimated). Corresponds to pde_system.ps. Set to SciMLBase.NullParameters() when there are no parameters.

• defaults: The default values of PDE parameters as a dictionary mapping each parameter symbol to its numeric value. Corresponds to pde_system.initial_conditions.

• default_p: The numeric default values of the PDE parameters as a plain Vector, extracted from defaults for use inside the loss function. nothing when eq_params is NullParameters.

• param_estim: Whether parameters are to be appended to the additional_loss

• additional_loss: The additional_loss function as provided by the user

• adaloss: The adaptive loss function

• depvars: The dependent variables of the system

• indvars: The independent variables of the system

• dict_indvars: A Dict{Symbol, Int} mapping each independent variable name (e.g. :x, :t) to its positional index in the coordinate vector. Used to build collocation point expressions.

• dict_depvars: A Dict{Symbol, Int} mapping each dependent variable name (e.g. :u, :v) to its positional index among the outputs. Used to index into phi and θ for multi-output systems.

• dict_depvar_input: A Dict{Symbol, Vector{Symbol}} mapping each dependent variable name to the list of independent variable names it depends on. For example, u(x, t) maps :u => [:x, :t]. Used to build the correct coordinate slices for each network input.

• logger: The logger as provided by the user

• multioutput: Whether there are multiple outputs, i.e. a system of PDEs

• iteration: The iteration counter used inside the cost function

• init_params: The initial parameters as provided by the user. If the PDE is a system of PDEs, this will be an array of arrays. If Lux.jl is used, then this is an array of ComponentArrays.

• flat_init_params: The initial parameters as a flattened array. This is the array that is used in the construction of the OptimizationProblem. If a Lux.jl neural network is used, then this flattened form is a ComponentArray. If the equation is a system of equations, then flat_init_params.depvar.x are the parameters for the neural network corresponding to the dependent variable x, and i.e. if depvar[i] == :x then for phi[i]. If param_estim = true, then flat_init_params.p are the parameters and flat_init_params.depvar.x are the neural network parameters, so flat_init_params.depvar.x would be the parameters of the neural network for the dependent variable x if it's a system.

• phi: The representation of the test function of the PDE solution

• derivative: The function used for computing the derivative

• strategy: The training strategy as provided by the user

• pde_indvars: For each PDE equation, the list of independent variables that appear in it. Used to build the correct collocation point layout for each loss term. For QuadratureTraining this holds the full argument list; for other strategies it holds only the variable symbols.

• bc_indvars: For each boundary condition equation, the list of independent variables that appear in it. Analogous to pde_indvars but for boundary loss terms.

• pde_integration_vars: For each PDE equation, the list of independent variables that are being integrated over (non-empty only when the equation contains a Symbolics.Integral term).

• bc_integration_vars: For each boundary condition equation, the list of independent variables that are being integrated over (non-empty only when the BC contains a Symbolics.Integral term).

• integral: The compiled numeric integral function, built by get_numeric_integral. Evaluates Symbolics.Integral terms at runtime using Integrals.jl quadrature.

• symbolic_pde_loss_functions: The PDE loss functions as represented in Julia AST

• symbolic_bc_loss_functions: The boundary condition loss functions as represented in Julia AST

• loss_functions: The PINNLossFunctions, i.e. the generated loss functions

PINNLossFunctions`

The generated functions from the PINNRepresentation

Fields

• bc_loss_functions: The boundary condition loss functions

• pde_loss_functions: The PDE loss functions

• full_loss_function: The full loss function, combining the PDE and boundary condition loss functions. This is the loss function that is used by the optimizer.

• additional_loss_function: The wrapped additional_loss, as pieced together for the optimizer.

• datafree_pde_loss_functions: The pre-data version of the PDE loss function

• datafree_bc_loss_functions: The pre-data version of the BC loss function

NNSDE(
    chain, opt, init_params = nothing; strategy = nothing, autodiff = false,
    batch = true, sub_batch = 1, strong_loss = false,
    moment_loss = false, param_estim = false, dataset = [],
    data_sub_batch = 1, numensemble = 10, additional_loss = nothing, kwargs...
)

This is an algorithm for solving stochastic ordinary differential equations using a specialization of physics-informed neural networks (PINNs). Allows users to solve standard SDEProblems using a Stochastic PINN (SPINN) solver.

Positional Arguments

• chain: A neural network (NN) architecture specific to SPINNs, such that the input dimensions correspond to time and n independent random variables chosen as the diffusion term's - Wiener process orthogonal random basis for it's KKL expansion. n has to be chosen by the user, depending on how accurately they want to represent stochasticity via the eigenvalues of the SDE's KKL expansion as in the SPINN loss function. The chain is defined as a Lux.AbstractLuxLayer or Flux.Chain. Flux.Chain will be converted to Lux using adapt(FromFluxAdaptor(), chain).
• opt: The optimizer to train the neural network.
• init_params: The initialization scheme for the neural network. By default, this is nothing which thus uses the random initialization provided by the neural network library.

Keyword Arguments

• strategy: The training strategy used to choose the points for the evaluations. Default of nothing means that QuadratureTraining with QuadGK is used if no dt is given, and GridTraining is used with dt if given. For the SDE solver, GridTraining is recommended for better weak solution estimates.

• autodiff: The switch between automatic and numerical differentiation for the PDE operators. The reverse mode of the loss function is always automatic differentiation (via Zygote), this is only for the derivative in the loss function (the derivative with respect to time).

• batch: The batch size for the loss computation. Defaults to true, means the neural network is applied at a row vector of values t simultaneously, i.e. it's the batch size for the neural network evaluations. This requires a neural network compatible with batched data. false means which means the application of the neural network is done at individual time points one at a time. This is not applicable to QuadratureTraining where batch is passed in the strategy which is the number of points it can parallelly compute the integrand.

• sub_batch: Interpretation depends on the training type chosen (based on strong_loss arg). In case of weak loss, training - this is the number of samples for each random coefficient z_i per timepoint to be taken, here a higher subbatch results in almost always a better capture of the weak solution. A defining feature of the loss and training strategy is that we construct training paths by taking separate, independent sets of `zifor each timepoint, and this is donen = sub_batchtimes for each timepoint. We essentially construct paths with Monte Carlo ensembles of random coefficients per time step, rather than continuous temporally consistent paths. In case of strong loss, training - this is the number of solution paths we are training over, the final SDEPINN solution is these fixed strong solution paths. Is1` by default.

• strong_loss: Controls the choice of training via the loss function aggregator operator and training type. If true, the loss has a strong form - summation across timepoints and n fixed (In each path, the same set of z_i coefficients for each timepoint) training solution paths (controlled via sub_batch), The solution returned is a strong solution for selected paths. If false, it is a weak loss - where n = sub_batch training solution paths are generated by random sampling across z_i probability space (In each path, independent random coeff z_i values for each timepoint). The solution returned is a weak solution for the SDE. Allows choosing weak/strong training discretization of the time and Gaussian coefficients domains. Note that weak training is almost always faster but if one is interested in pathwise solutions of the SDE then strong_loss must be true. It is false by default.

• moment_loss: Allows user to include a moment matching loss (1st and 2nd - mean and var) for the solver. It is calculated against the dataset provided. Is false by default.

• param_estim: Boolean to indicate whether parameters of the differential equations are learnt along with parameters of the neural network.

• dataset: A dataset used to train the SDEPINN using observed process samples, at respective timepoints. The L2 loss, moment_loss is created using this. It is a Vector of x, t where x is a nested vector of multiple observations of the adapted process being learnt, here each inner vector corresponds to one strong solution set of timeseries observations. t is a vector of timepoints at which we have the multiple x observations.

• data_sub_batch: The number of sets of random coefficients zi to be taken for each timepoint while matching moments in `momentloss. The moment matching is done against strong process observations, therefore to match moments we must have corresponding SDEPINN outputs and inputs. Since realistically we cant get the eigenfunction, random coefficient decompositions for each observation we try to approximate the final moments by taking mean/sum over multiple sets of z_i concatenated with timepoints as SDEPINN inputs to get outputs that we can finally match with the dataset mean and variance. Is internallymax(datasubbatch, length(dataset[1]))` by default.

• numensemble: The solver returns an ensemble results/weak solution over user's provided saveat discretization. numensemble controls the Number of solution predictions to take an ensemble over for each timepoints. Is 10 by default.

• additional_loss: A function additional_loss(phi, θ) where phi are the neural network trial solutions, θ are the weights of the neural network(s).

• kwargs: Extra keyword arguments are splatted to the Optimization.jl solve call.

Examples

u0 = [1.0, 1.0]
ts = [t for t in 1:100]
(u_, t_) = (analytical_func(ts), ts)
function additional_loss(phi, θ)
    return sum(sum(abs2, [phi(t, θ) for t in t_] .- u_)) / length(u_)
end
alg = NNSDE(chain, opt, additional_loss = additional_loss)

u₀ = 0.5
f(u, p, t) = 1.2 * u
g(u, p, t) = 1.1 * u
tspan = (0.0, 1.0)
prob = SDEProblem(f, g, u₀, tspan)
n_z = 3
dim = 1 + n_z
luxchain = Chain(Dense(dim, 16, σ), Dense(16, 16, σ), Dense(16, 1))
opt = BFGS()
sol = solve(prob, NNSDE(luxchain, opt), verbose = true, dt = 1 / 50.0f0, abstol = 1.0e-10, maxiters = 200)

Solution Notes

Note that the returned weak solution is evaluated at fixed time points according to standard output handlers such as saveat and dt. However, the neural network is a fully continuous solution so sol(t) is an accurate interpolation (up to the neural network training result). In addition, the OptimizationSolution is returned as sol.k for further analysis.

References

Stochastic Physics-Informed Neural Ordinary Differential Equations : https://arxiv.org/abs/2109.01621 Stochastic PDE Functionality #531 : https://github.com/SciML/NeuralPDE.jl/issues/531

SDEPINN(;
    chain, x_0, x_end, optimalg = nothing, norm_loss_alg = nothing,
    initial_parameters = nothing, Nt = 20, dx = 0.05, σ_var_bc = 0.05,
    λ_ic = 1.0, λ_norm = 1.0, distrib = Normal(0.5, 0.01), strategy = nothing,
    autodiff = true, batch = false, param_estim = false, dataset = nothing,
    additional_loss = nothing, kwargs...
)

Solve an SDEProblem by training a physics-informed neural network on its associated Fokker-Planck equation over the spatial interval from x_0 to x_end.

Keyword Arguments

• chain: Neural network used to represent the probability density.
• x_0, x_end: Lower and upper endpoints of the spatial domain.
• optimalg: Optimizer used to train the network.
• norm_loss_alg: Integration algorithm used by the normalization loss.
• initial_parameters: Initial network parameters. NeuralPDE initializes them when omitted.
• Nt: Number of temporal training points.
• dx: Spatial grid spacing used to evaluate the solution.
• σ_var_bc: Width of the Gaussian approximation to the initial condition.
• λ_ic, λ_norm: Initial-condition and normalization loss weights.
• distrib: Initial probability distribution.
• strategy: Optional training strategy.
• autodiff: Whether to use automatic differentiation for temporal derivatives.
• batch: Whether to batch training points.
• param_estim: Whether to estimate equation parameters with the network parameters.
• dataset: Optional observed data used during training.
• additional_loss: Optional function that contributes an additional training loss.
• kwargs: Additional keyword arguments forwarded to the Optimization.jl solve.

Example

using Integrals, Lux, NeuralPDE, OptimizationOptimJL

chain = Chain(Dense(2, 16, tanh), Dense(16, 1))
alg = SDEPINN(
    chain = chain,
    optimalg = BFGS(),
    norm_loss_alg = HCubatureJL(),
    x_0 = -4.0,
    x_end = 4.0,
)