SSE not MSE, PDE solver

src/PDE_BPINN.jl

@@ -53,18 +53,17 @@ function get_lossy(pinnrep, dataset, Dict_differentials)
     colloc_equations = [Symbolics.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)
+    # nested vector of data_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)
+    data_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
+    return data_colloc_loss_functions
 end
 
 function get_symbols(dataset, depvars, eqs)
@@ -321,39 +320,38 @@ function ahmc_bayesian_pinn_pde(pde_system, discretization;
     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
+        data_colloc_loss_functions = get_lossy(pinnrep, dataset_pde, Dict_differentials)
+        # size = 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)]
+        # j is number of indvar coords in dataset, i is number of PDE equations in system
+        # -1 is placeholder, removed in merge_strategy_with_loglikelihood_function function call (train_sets[:, 2:end]())
+        colloc_train_sets = [[hcat([-1], train_sets_pde[i][:, j]...)
+                              for i in eachindex(data_colloc_loss_functions[1])]
+                             for j in eachindex(data_colloc_loss_functions)]
 
-        # for each datafree_colloc_loss_function create loss_functions by passing dataset's indvar coords as train_sets_pde.
+        # using dataset's indvar coords as train_sets_pde and indvar coord's datafree_colloc_loss_function, create loss functions
         # 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],
+                                        data_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)]
+                                    for i in eachindex(data_colloc_loss_functions)]
 
         function L2_loss2(θ, phynewstd)
-            # first vector of losses,from tuple -> pde losses, first[1] pde loss
-            pde_loglikelihoods = [sum([pde_loss_function(θ, phynewstd[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(θ, phynewstd[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)
+            # first sum is over points losses over many equations for the same points
+            # second sum is over all points
+            pde_loglikelihoods = sum([sum([pde_loss_function(θ,
+                                               phynewstd[i])
+                                           for (i, pde_loss_function) in enumerate(pde_loss_functions)])
+                                      for pde_loss_functions in pde_loss_function_points])
         end
     end
 
@@ -435,7 +433,7 @@ function ahmc_bayesian_pinn_pde(pde_system, discretization;
         @printf("Current Physics Log-likelihood : %g\n",
             ℓπ.full_loglikelihood(setparameters(ℓπ, initial_θ), ℓπ.allstd))
         @printf("Current Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, initial_θ))
-        @printf("Current MSE against dataset Log-likelihood : %g\n",
+        @printf("Current SSE against dataset Log-likelihood : %g\n",
             L2LossData(ℓπ, initial_θ))
         if !(newloss isa Nothing)
             @printf("Current new loss : %g\n",
@@ -493,7 +491,7 @@ function ahmc_bayesian_pinn_pde(pde_system, discretization;
             @printf("Final Physics Log-likelihood : %g\n",
                 ℓπ.full_loglikelihood(setparameters(ℓπ, samples[end]), ℓπ.allstd))
             @printf("Final Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, samples[end]))
-            @printf("Final MSE against dataset Log-likelihood : %g\n",
+            @printf("Final SSE against dataset Log-likelihood : %g\n",
                 L2LossData(ℓπ, samples[end]))
             if !(newloss isa Nothing)
                 @printf("Final new loss : %g\n",

src/advancedHMC_MCMC.jl

@@ -425,7 +425,7 @@ function ahmc_bayesian_pinn_ode(
     if verbose
         @printf("Current Physics Log-likelihood: %g\n", physloglikelihood(ℓπ, initial_θ))
         @printf("Current Prior Log-likelihood: %g\n", priorweights(ℓπ, initial_θ))
-        @printf("Current MSE against dataset Log-likelihood: %g\n",
+        @printf("Current SSE against dataset Log-likelihood: %g\n",
             L2LossData(ℓπ, initial_θ))
         if estim_collocate
             @printf("Current gradient loss against dataset Log-likelihood: %g\n",
@@ -487,7 +487,7 @@ function ahmc_bayesian_pinn_ode(
             @printf("Final Physics Log-likelihood: %g\n",
                 physloglikelihood(ℓπ, samples[end]))
             @printf("Final Prior Log-likelihood: %g\n", priorweights(ℓπ, samples[end]))
-            @printf("Final MSE against dataset Log-likelihood: %g\n",
+            @printf("Final SSE against dataset Log-likelihood: %g\n",
                 L2LossData(ℓπ, samples[end]))
             if estim_collocate
                 @printf("Final gradient loss against dataset Log-likelihood: %g\n",

src/discretize.jl

@@ -526,6 +526,10 @@ function SciMLBase.symbolic_discretize(pde_system::PDESystem, discretization::Ab
     function get_likelihood_estimate_function(discretization::BayesianPINN)
         dataset_pde, dataset_bc = discretization.dataset
 
+        pde_loss_functions, bc_loss_functions = merge_strategy_with_loglikelihood_function(
+            pinnrep, strategy,
+            datafree_pde_loss_functions, datafree_bc_loss_functions)
+
         # required as Physics loss also needed on the discrete dataset domain points
         # data points are discrete and so by default GridTraining loss applies
         # passing placeholder dx with GridTraining, it uses data points irl
@@ -542,20 +546,22 @@ function SciMLBase.symbolic_discretize(pde_system::PDESystem, discretization::Ab
         function full_loss_function(θ, allstd::Vector{Vector{Float64}})
             stdpdes, stdbcs, stdextra = allstd
             # the aggregation happens on cpu even if the losses are gpu, probably fine since it's only a few of them
-            pde_loglikelihoods = [logpdf(Normal(0, stdpdes[i]), pde_loss_function(θ))
-                                  for (i, pde_loss_function) in enumerate(pde_loss_functions)]
+            # SSE FOR LOSS ON GRIDPOINTS not MSE ! i, j depend on number of bcs and eqs
+            pde_loglikelihoods = sum([pde_loglike_function(θ, stdpdes[i])
+                                      for (i, pde_loglike_function) in enumerate(pde_loss_functions)])
 
-            bc_loglikelihoods = [logpdf(Normal(0, stdbcs[j]), bc_loss_function(θ))
-                                 for (j, bc_loss_function) in enumerate(bc_loss_functions)]
+            bc_loglikelihoods = sum([bc_loglike_function(θ, stdbcs[j])
+                                     for (j, bc_loglike_function) in enumerate(bc_loss_functions)])
 
+            # final newloss creation components are similar to this
             if !(datapde_loss_functions isa Nothing)
-                pde_loglikelihoods += [pde_loglike_function(θ, allstd[1])
-                                       for (j, pde_loglike_function) in enumerate(datapde_loss_functions)]
+                pde_loglikelihoods += sum([pde_loglike_function(θ, allstd[1])
+                                           for (j, pde_loglike_function) in enumerate(datapde_loss_functions)])
             end
 
             if !(databc_loss_functions isa Nothing)
-                bc_loglikelihoods += [bc_loglike_function(θ, allstd[2])
-                                      for (j, bc_loglike_function) in enumerate(databc_loss_functions)]
+                bc_loglikelihoods += sum([bc_loglike_function(θ, allstd[2])
+                                          for (j, bc_loglike_function) in enumerate(databc_loss_functions)])
             end
 
             # this is kind of a hack, and means that whenever the outer function is evaluated the increment goes up, even if it's not being optimized

src/training_strategies.jl

@@ -55,23 +55,39 @@ function merge_strategy_with_loglikelihood_function(pinnrep::PINNRepresentation,
 
     # only when physics loss is taken, merge_strategy_with_loglikelihood_function() call case
     if ((train_sets_bc isa Nothing) && (train_sets_pde isa Nothing))
-        train_sets_pde, train_sets_bc = generate_training_sets(
-            domains, dx, eqs, bcs, eltypeθ,
-            dict_indvars, dict_depvars)
+        train_sets = generate_training_sets(
+            pinnrep.domains, strategy.dx, pinnrep.eqs, pinnrep.bcs, eltypeθ,
+            pinnrep.dict_indvars, pinnrep.dict_depvars)
+
+        train_sets_pde, train_sets_bc = train_sets |> adaptor
+        # train_sets_pde matches PhysicsInformedNN solver train_sets[1] dims.
+        pde_loss_functions = [get_points_loss_functions(_loss, _set, eltypeθ, strategy)
+                              for (_loss, _set) in zip(
+            datafree_pde_loss_function, train_sets_pde)]
+
+        bc_loss_functions = [get_points_loss_functions(_loss, _set, eltypeθ, strategy)
+                             for (_loss, _set) in zip(
+            datafree_bc_loss_function, train_sets_bc)]
+
+        return pde_loss_functions, bc_loss_functions
     end
 
-    # is vec as later each _set in pde_train_sets are columns as points transformed to
-    # vector of points (pde_train_sets must be rowwise)
     pde_loss_functions = if train_sets_pde !== nothing
-        pde_train_sets = [train_set[:, 2:end] for train_set in train_sets_pde] |> adaptor
+        # as first col in all rows is depvar's value in depvar's dataset respectively
+        # and we want only all depvar dataset's indvar points
+        pde_train_sets = [train_set[:, 2:end]' for train_set in train_sets_pde] |> adaptor
+
+        # pde_train_sets must match PhysicsInformedNN solver train_sets[1] dims. It is a vector with coords.
+        # Vector is for number of PDE equations in system, Matrix has rows of indvar grid point coords
+        # each loss struct mapped onto (total_numpoints_combs, dim_indvars)
         [get_points_loss_functions(_loss, _set, eltypeθ, strategy)
          for (_loss, _set) in zip(datafree_pde_loss_function, pde_train_sets)]
     else
         nothing
     end
 
     bc_loss_functions = if train_sets_bc !== nothing
-        bcs_train_sets = [train_set[:, 2:end] for train_set in train_sets_bc] |> adaptor
+        bcs_train_sets = [train_set[:, 2:end]' for train_set in train_sets_bc] |> adaptor
         [get_points_loss_functions(_loss, _set, eltypeθ, strategy)
          for (_loss, _set) in zip(datafree_bc_loss_function, bcs_train_sets)]
     else

test/BPINN_PDE_tests.jl

@@ -479,7 +479,7 @@ end
     sol_new = ahmc_bayesian_pinn_pde(pde_system,
         discretization;
         draw_samples = 150,
-        bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], phystdnew = [0.2],
+        bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], phynewstd = [0.2],
         phystd = [0.2], l2std = [0.5], param = [Distributions.Normal(2.0, 2)],
         priorsNNw = (0.0, 1.0),
         saveats = [1 / 100.0, 1 / 100.0],