better BPINN formulation, improvements to include dataset domain points

src/BPINN_ode.jl

@@ -100,6 +100,8 @@ struct BNNODE{C, K, IT <: NamedTuple,
     init_params::I
     Adaptorkwargs::A
     Integratorkwargs::IT
+    numensemble::Int64
+    estim_collocate::Bool
     autodiff::Bool
     progress::Bool
     verbose::Bool
@@ -112,6 +114,8 @@ function BNNODE(chain, Kernel = HMC; strategy = nothing, draw_samples = 2000,
             Metric = DiagEuclideanMetric,
             targetacceptancerate = 0.8),
         Integratorkwargs = (Integrator = Leapfrog,),
+    numensemble = floor(Int, draw_samples / 3),
+    estim_collocate = false,
         autodiff = false, progress = false, verbose = false)
     !(chain isa Lux.AbstractExplicitLayer) &&
         (chain = adapt(FromFluxAdaptor(false, false), chain))
@@ -120,6 +124,7 @@ function BNNODE(chain, Kernel = HMC; strategy = nothing, draw_samples = 2000,
         phystd, dataset, physdt, MCMCkwargs,
         nchains, init_params,
         Adaptorkwargs, Integratorkwargs,
+        numensemble, estim_collocate,
         autodiff, progress, verbose)
 end
 
@@ -186,7 +191,7 @@ function SciMLBase.__solve(prob::SciMLBase.ODEProblem,
     @unpack chain, l2std, phystd, param, priorsNNw, Kernel, strategy,
     draw_samples, dataset, init_params,
     nchains, physdt, Adaptorkwargs, Integratorkwargs,
-    MCMCkwargs, autodiff, progress, verbose = alg
+    MCMCkwargs, numensemble, estim_collocate, autodiff, progress, verbose = alg
 
     # ahmc_bayesian_pinn_ode needs param=[] for easier vcat operation for full vector of parameters
     param = param === nothing ? [] : param
@@ -211,7 +216,8 @@ function SciMLBase.__solve(prob::SciMLBase.ODEProblem,
         Integratorkwargs = Integratorkwargs,
         MCMCkwargs = MCMCkwargs,
         progress = progress,
-        verbose = verbose)
+        verbose = verbose,
+        estim_collocate = estim_collocate)
 
     fullsolution = BPINNstats(mcmcchain, samples, statistics)
     ninv = length(param)
@@ -220,7 +226,8 @@ function SciMLBase.__solve(prob::SciMLBase.ODEProblem,
     if chain isa Lux.AbstractExplicitLayer
         θinit, st = Lux.setup(Random.default_rng(), chain)
         θ = [vector_to_parameters(samples[i][1:(end - ninv)], θinit)
-             for i in (draw_samples - numensemble):draw_samples]
+             for i in 1:max(draw_samples - draw_samples ÷ 10, draw_samples - 1000)]
+
         luxar = [chain(t', θ[i], st)[1] for i in 1:numensemble]
         # only need for size
         θinit = collect(ComponentArrays.ComponentArray(θinit))

src/PDE_BPINN.jl

@@ -4,6 +4,7 @@ mutable struct PDELogTargetDensity{
     P <: Vector{<:Distribution},
     I,
     F,
+    FF,
     PH
 }
     dim::Int64
@@ -15,17 +16,19 @@ mutable struct PDELogTargetDensity{
     extraparams::Int
     init_params::I
     full_loglikelihood::F
+    L2_loss2::FF
     Φ::PH
 
     function PDELogTargetDensity(dim, strategy, dataset,
             priors, allstd, names, extraparams,
-            init_params::AbstractVector, full_loglikelihood, Φ)
+            init_params::AbstractVector, full_loglikelihood, L2_loss2, Φ)
         new{
             typeof(strategy),
             typeof(dataset),
             typeof(priors),
             typeof(init_params),
             typeof(full_loglikelihood),
+            typeof(L2_loss2),
             typeof(Φ)
         }(dim,
             strategy,
@@ -36,18 +39,20 @@ mutable struct PDELogTargetDensity{
             extraparams,
             init_params,
             full_loglikelihood,
+            L2_loss2,
             Φ)
     end
     function PDELogTargetDensity(dim, strategy, dataset,
             priors, allstd, names, extraparams,
             init_params::Union{NamedTuple, ComponentArrays.ComponentVector},
-            full_loglikelihood, Φ)
+            full_loglikelihood, L2_loss2, Φ)
         new{
             typeof(strategy),
             typeof(dataset),
             typeof(priors),
             typeof(init_params),
             typeof(full_loglikelihood),
+            typeof(L2_loss2),
             typeof(Φ)
         }(dim,
             strategy,
@@ -58,22 +63,87 @@ mutable struct PDELogTargetDensity{
             extraparams,
             init_params,
             full_loglikelihood,
+            L2_loss2,
             Φ)
     end
 end
 
+# you get a vector of losses
+function get_lossy(pinnrep, dataset, Dict_differentials)
+    eqs = pinnrep.eqs
+    depvars = pinnrep.depvars #depvar order is same as dataset
+
+    # Dict_differentials is filled with Differential operator => diff_i key-value pairs
+    # masking operation
+    eqs_new = substitute.(eqs, Ref(Dict_differentials))
+
+    to_subs, tobe_subs = get_symbols(dataset, depvars, eqs)
+
+    # for values of all depvars at corresponding indvar values in dataset, create dictionaries {Dict(x(t) => 1.0496435863173237, y(t) => 1.9227770685615337)}
+    # In each Dict, num form of depvar is key to its value at certain coords of indvars, n_dicts = n_rows_dataset(or n_indvar_coords_dataset)
+    eq_subs = [Dict(tobe_subs[depvar] => to_subs[depvar][i] for depvar in depvars)
+               for i in 1:size(dataset[1][:, 1])[1]]
+
+    # for each dataset point(eq_sub dictionary), substitute in masked equations
+    # n_collocated_equations = n_rows_dataset(or n_indvar_coords_dataset)
+    masked_colloc_equations = [[substitute(eq, eq_sub) for eq in eqs_new]
+                            for eq_sub in eq_subs]
+    # now we have vector of dataset depvar's collocated equations
+
+    # reverse dict for re-substituting values of Differential(t)(u(t)) etc
+    rev_Dict_differentials = Dict(value => key for (key, value) in Dict_differentials)
+
+    # unmask Differential terms in masked_colloc_equations
+    colloc_equations = [substitute.(masked_colloc_equation, Ref(rev_Dict_differentials))
+                    for masked_colloc_equation in masked_colloc_equations]
+
+    # nested vector of datafree_pde_loss_functions (as in discretize.jl)
+    # each sub vector has dataset's indvar coord's datafree_colloc_loss_function, n_subvectors = n_rows_dataset(or n_indvar_coords_dataset)
+    # zip each colloc equation with args for each build_loss call per equation vector
+    datafree_colloc_loss_functions = [[build_loss_function(pinnrep, eq, pde_indvar)
+                       for (eq, pde_indvar, integration_indvar) in zip(colloc_equation,
+        pinnrep.pde_indvars,
+        pinnrep.pde_integration_vars)] for colloc_equation in colloc_equations]
+
+    return datafree_colloc_loss_functions
+end
+
+function get_symbols(dataset, depvars, eqs)
+    # take only values of depvars from dataset
+    depvar_vals = [dataset_i[:, 1] for dataset_i in dataset]
+    # order of pinnrep.depvars, depvar_vals, BayesianPINN.dataset must be same
+    to_subs = Dict(depvars .=> depvar_vals)
+
+    numform_vars = Symbolics.get_variables.(eqs)
+    Eq_vars = unique(reduce(vcat, numform_vars))
+    # got equation's depvar num format {x(t)} for use in substitute()
+
+    tobe_subs = Dict()
+    for a in depvars
+        for i in Eq_vars
+            expr = toexpr(i)
+            if (expr isa Expr) && (expr.args[1] == a)
+                tobe_subs[a] = i
+            end
+        end
+    end
+    # depvar symbolic and num format got, tobe_subs : Dict{Any, Any}(:y => y(t), :x => x(t))
+
+    return to_subs, tobe_subs
+end
+
 function LogDensityProblems.logdensity(Tar::PDELogTargetDensity, θ)
     # for parameter estimation neccesarry to use multioutput case
-    return Tar.full_loglikelihood(setparameters(Tar, θ),
-               Tar.allstd) + priorlogpdf(Tar, θ) + L2LossData(Tar, θ)
-    # + L2loss2(Tar, θ)
+    if Tar.L2_loss2 isa Nothing
+        return Tar.full_loglikelihood(setparameters(Tar, θ), Tar.allstd) +
+               priorlogpdf(Tar, θ) + L2LossData(Tar, θ)
+    else
+        return Tar.full_loglikelihood(setparameters(Tar, θ), Tar.allstd) +
+               priorlogpdf(Tar, θ) + L2LossData(Tar, θ) +
+               Tar.L2_loss2(setparameters(Tar, θ), Tar.allstd)
+    end
 end
 
-# function L2loss2(Tar::PDELogTargetDensity, θ)
-#     return Tar.full_loglikelihood(setparameters(Tar, θ),
-#         Tar.allstd)
-# end
-
 function setparameters(Tar::PDELogTargetDensity, θ)
     names = Tar.names
     ps_new = θ[1:(end - Tar.extraparams)]
@@ -117,6 +187,8 @@ function L2LossData(Tar::PDELogTargetDensity, θ)
 
     # dataset of form Vector[matrix_x, matrix_y, matrix_z]
     # matrix_i is of form [i,indvar1,indvar2,..] (needed in case if heterogenous domains)
+    # note that indvar1,indvar2.. cols can be different values for different depvar matrices
+    # dataset,phi order follows pinnrep.depvars orders of variables (order of declaration in @variables macro)
 
     # Phi is the trial solution for each NN in chain array
     # Creating logpdf( MvNormal(Phi(t,θ),std), dataset[i] )
@@ -162,27 +234,6 @@ function priorlogpdf(Tar::PDELogTargetDensity, θ)
     return logpdf(nnwparams, θ)
 end
 
-function integratorchoice(Integratorkwargs, initial_ϵ)
-    Integrator = Integratorkwargs[:Integrator]
-    if Integrator == JitteredLeapfrog
-        jitter_rate = Integratorkwargs[:jitter_rate]
-        Integrator(initial_ϵ, jitter_rate)
-    elseif Integrator == TemperedLeapfrog
-        tempering_rate = Integratorkwargs[:tempering_rate]
-        Integrator(initial_ϵ, tempering_rate)
-    else
-        Integrator(initial_ϵ)
-    end
-end
-
-function adaptorchoice(Adaptor, mma, ssa)
-    if Adaptor != AdvancedHMC.NoAdaptation()
-        Adaptor(mma, ssa)
-    else
-        AdvancedHMC.NoAdaptation()
-    end
-end
-
 function inference(samples, pinnrep, saveats, numensemble, ℓπ)
     domains = pinnrep.domains
     phi = pinnrep.phi
@@ -247,6 +298,27 @@ function inference(samples, pinnrep, saveats, numensemble, ℓπ)
     return ensemblecurves, estimatedLuxparams, estimated_params, timepoints
 end
 
+function integratorchoice(Integratorkwargs, initial_ϵ)
+    Integrator = Integratorkwargs[:Integrator]
+    if Integrator == JitteredLeapfrog
+        jitter_rate = Integratorkwargs[:jitter_rate]
+        Integrator(initial_ϵ, jitter_rate)
+    elseif Integrator == TemperedLeapfrog
+        tempering_rate = Integratorkwargs[:tempering_rate]
+        Integrator(initial_ϵ, tempering_rate)
+    else
+        Integrator(initial_ϵ)
+    end
+end
+
+function adaptorchoice(Adaptor, mma, ssa)
+    if Adaptor != AdvancedHMC.NoAdaptation()
+        Adaptor(mma, ssa)
+    else
+        AdvancedHMC.NoAdaptation()
+    end
+end
+
 """
     ahmc_bayesian_pinn_pde(pde_system, discretization;
             draw_samples = 1000,
@@ -295,15 +367,56 @@ end
 function ahmc_bayesian_pinn_pde(pde_system, discretization;
         draw_samples = 1000,
         bcstd = [0.01], l2std = [0.05],
-        phystd = [0.05], priorsNNw = (0.0, 2.0),
+        phystd = [0.05], phystdnew = [0.05], priorsNNw = (0.0, 2.0),
         param = [], nchains = 1, Kernel = HMC(0.1, 30),
         Adaptorkwargs = (Adaptor = StanHMCAdaptor,
             Metric = DiagEuclideanMetric, targetacceptancerate = 0.8),
         Integratorkwargs = (Integrator = Leapfrog,), saveats = [1 / 10.0],
-        numensemble = floor(Int, draw_samples / 3), progress = false, verbose = false)
+        numensemble = floor(Int, draw_samples / 3), Dict_differentials = nothing,
+        progress = false, verbose = false)
     pinnrep = symbolic_discretize(pde_system, discretization)
     dataset_pde, dataset_bc = discretization.dataset
 
+    newloss = if Dict_differentials isa Nothing
+        nothing
+    else
+        datafree_colloc_loss_functions = get_lossy(pinnrep, dataset_pde, Dict_differentials)
+        # equals number of indvar coords in dataset
+        # add case for if parameters present in bcs?
+
+        train_sets_pde = get_dataset_train_points(pde_system.eqs,
+                dataset_pde,
+                pinnrep)
+        colloc_train_sets = [[hcat(train_sets_pde[i][:, j]...)' for i in eachindex(datafree_colloc_loss_functions[1])] for j in eachindex(datafree_colloc_loss_functions)]
+
+        # for each datafree_colloc_loss_function create loss_functions by passing dataset's indvar coords as train_sets_pde.
+        # placeholder strategy = GridTraining(0.1), datafree_bc_loss_function and train_sets_bc must be nothing
+        # order of indvar coords will be same as corresponding depvar coords values in dataset provided in get_lossy() call.
+        pde_loss_function_points = [merge_strategy_with_loglikelihood_function(
+            pinnrep,
+            GridTraining(0.1),
+            datafree_colloc_loss_functions[i],
+            nothing;
+            train_sets_pde = colloc_train_sets[i],
+            train_sets_bc = nothing)[1]
+                              for i in eachindex(datafree_colloc_loss_functions)]
+
+        function L2_loss2(θ, allstd)
+            stdpdesnew = allstd[4]
+
+            # first vector of losses,from tuple -> pde losses, first[1] pde loss
+            pde_loglikelihoods = [sum([pde_loss_function(θ, stdpdesnew[i])
+                                       for (i, pde_loss_function) in enumerate(pde_loss_functions)])
+                                  for pde_loss_functions in pde_loss_function_points]
+
+            # bc_loglikelihoods = [sum([bc_loss_function(θ, stdpdesnew[i]) for (i, bc_loss_function) in enumerate(pde_loss_function_points[1])]) for pde_loss_function_points in pde_loss_functions]
+            #                      for (j, bc_loss_function) in enumerate(bc_loss_functions)]
+
+            return sum(pde_loglikelihoods)
+        end
+    end
+
+    # [WIP] add overall functionality for BC dataset points
     if ((dataset_bc isa Nothing) && (dataset_pde isa Nothing))
         dataset = nothing
     elseif dataset_bc isa Nothing
@@ -333,9 +446,6 @@ function ahmc_bayesian_pinn_pde(pde_system, discretization;
     # NN solutions for loglikelihood which is used for L2lossdata
     Φ = pinnrep.phi
 
-    # for new L2 loss
-    # discretization.additional_loss = 
-
     if nchains < 1
         throw(error("number of chains must be greater than or equal to 1"))
     end
@@ -375,11 +485,12 @@ function ahmc_bayesian_pinn_pde(pde_system, discretization;
         strategy,
         dataset,
         priors,
-        [phystd, bcstd, l2std],
+        [phystd, bcstd, l2std, phystdnew],
         names,
         ninv,
         initial_nnθ,
         full_weighted_loglikelihood,
+        newloss,
         Φ)
 
     Adaptor, Metric, targetacceptancerate = Adaptorkwargs[:Adaptor],
@@ -394,10 +505,14 @@ function ahmc_bayesian_pinn_pde(pde_system, discretization;
             ℓπ.allstd))
     @info("Current Prior Log-likelihood : ", priorlogpdf(ℓπ, initial_θ))
     @info("Current MSE against dataset Log-likelihood : ", L2LossData(ℓπ, initial_θ))
+    if !(newloss isa Nothing)
+        @info("Current L2_LOSSY : ",
+            ℓπ.L2_loss2(setparameters(ℓπ, initial_θ),
+                ℓπ.allstd))
+    end
 
     # parallel sampling option
     if nchains != 1
-
         # Cache to store the chains
         bpinnsols = Vector{Any}(undef, nchains)
 
@@ -453,6 +568,11 @@ function ahmc_bayesian_pinn_pde(pde_system, discretization;
         @info("Current Prior Log-likelihood : ", priorlogpdf(ℓπ, samples[end]))
         @info("Current MSE against dataset Log-likelihood : ",
             L2LossData(ℓπ, samples[end]))
+        if !(newloss isa Nothing)
+            @info("Current L2_LOSSY : ",
+                ℓπ.L2_loss2(setparameters(ℓπ, samples[end]),
+                    ℓπ.allstd))
+        end
 
         fullsolution = BPINNstats(mcmc_chain, samples, stats)
         ensemblecurves, estimnnparams, estimated_params, timepoints = inference(samples,