diff --git a/src/BPINN_ode.jl b/src/BPINN_ode.jl
index f65f1d659..489cd5f6d 100644
--- a/src/BPINN_ode.jl
+++ b/src/BPINN_ode.jl
@@ -3,7 +3,7 @@
 """
     BNNODE(chain, kernel = HMC; strategy = nothing, draw_samples = 2000,
            priorsNNw = (0.0, 2.0), param = [nothing], l2std = [0.05],
-           phystd = [0.05], dataset = [nothing], physdt = 1 / 20.0,
+           phystd = [0.05], phynewstd = [0.05], dataset = [nothing], physdt = 1 / 20.0,
            MCMCargs = (; n_leapfrog=30), nchains = 1, init_params = nothing,
            Adaptorkwargs = (; Adaptor = StanHMCAdaptor, targetacceptancerate = 0.8,
                               Metric = DiagEuclideanMetric),
@@ -86,6 +86,7 @@ Kevin Linka, Amelie Schäfer, Xuhui Meng, Zongren Zou, George Em Karniadakis, El
     param <: Union{Nothing, Vector{<:Distribution}}
     l2std::Vector{Float64}
     phystd::Vector{Float64}
+    phynewstd::Vector{Float64}
     dataset <: Union{Vector{Nothing}, Vector{<:Vector{<:AbstractFloat}}}
     physdt::Float64
     MCMCkwargs <: NamedTuple
@@ -100,10 +101,10 @@ Kevin Linka, Amelie Schäfer, Xuhui Meng, Zongren Zou, George Em Karniadakis, El
     verbose::Bool
 end
 
-function BNNODE(chain, kernel = HMC; strategy = nothing, draw_samples = 2000,
+function BNNODE(chain, kernel = HMC; strategy = nothing, draw_samples = 1000,
         priorsNNw = (0.0, 2.0), param = nothing, l2std = [0.05], phystd = [0.05],
-        dataset = [nothing], physdt = 1 / 20.0, MCMCkwargs = (n_leapfrog = 30,),
-        nchains = 1, init_params = nothing,
+        phynewstd = [0.05], dataset = [nothing], physdt = 1 / 20.0,
+        MCMCkwargs = (n_leapfrog = 30,), nchains = 1, init_params = nothing,
         Adaptorkwargs = (Adaptor = StanHMCAdaptor,
             Metric = DiagEuclideanMetric, targetacceptancerate = 0.8),
         Integratorkwargs = (Integrator = Leapfrog,),
@@ -111,7 +112,7 @@ function BNNODE(chain, kernel = HMC; strategy = nothing, draw_samples = 2000,
         estim_collocate = false, autodiff = false, progress = false, verbose = false)
     chain isa AbstractLuxLayer || (chain = FromFluxAdaptor()(chain))
     return BNNODE(chain, kernel, strategy, draw_samples, priorsNNw, param, l2std, phystd,
-        dataset, physdt, MCMCkwargs, nchains, init_params, Adaptorkwargs,
+        phynewstd, dataset, physdt, MCMCkwargs, nchains, init_params, Adaptorkwargs,
         Integratorkwargs, numensemble, estim_collocate, autodiff, progress, verbose)
 end
 
@@ -157,7 +158,7 @@ end
 function SciMLBase.__solve(prob::SciMLBase.ODEProblem, alg::BNNODE, args...; dt = nothing,
         timeseries_errors = true, save_everystep = true, adaptive = false,
         abstol = 1.0f-6, reltol = 1.0f-3, verbose = false, saveat = 1 / 50.0,
-        maxiters = nothing, numensemble = floor(Int, alg.draw_samples / 3))
+        maxiters = nothing)
     (; chain, param, strategy, draw_samples, numensemble, verbose) = alg
 
     # ahmc_bayesian_pinn_ode needs param=[] for easier vcat operation for full vector of parameters
@@ -168,7 +169,8 @@ function SciMLBase.__solve(prob::SciMLBase.ODEProblem, alg::BNNODE, args...; dt
 
     mcmcchain, samples, statistics = ahmc_bayesian_pinn_ode(
         prob, chain; strategy, alg.dataset, alg.draw_samples, alg.init_params,
-        alg.physdt, alg.l2std, alg.phystd, alg.priorsNNw, param, alg.nchains, alg.autodiff,
+        alg.physdt, alg.l2std, alg.phystd, alg.phynewstd,
+        alg.priorsNNw, param, alg.nchains, alg.autodiff,
         Kernel = alg.kernel, alg.Adaptorkwargs, alg.Integratorkwargs,
         alg.MCMCkwargs, alg.progress, alg.verbose, alg.estim_collocate)
 
@@ -178,7 +180,7 @@ function SciMLBase.__solve(prob::SciMLBase.ODEProblem, alg::BNNODE, args...; dt
 
     θinit, st = LuxCore.setup(Random.default_rng(), chain)
     θ = [vector_to_parameters(samples[i][1:(end - ninv)], θinit)
-         for i in 1:max(draw_samples - draw_samples ÷ 10, draw_samples - 1000)]
+         for i in (draw_samples - numensemble):draw_samples]
 
     luxar = [chain(t', θ[i], st)[1] for i in 1:numensemble]
     # only need for size
diff --git a/src/PDE_BPINN.jl b/src/PDE_BPINN.jl
index c57bcd71c..3ea00758f 100644
--- a/src/PDE_BPINN.jl
+++ b/src/PDE_BPINN.jl
@@ -4,17 +4,90 @@
     dataset <: Union{Nothing, Vector{<:Matrix{<:Real}}}
     priors <: Vector{<:Distribution}
     allstd::Vector{Vector{Float64}}
+    phynewstd::Vector{Float64}
     names::Tuple
     extraparams::Int
     init_params <: Union{AbstractVector, NamedTuple, ComponentArray}
     full_loglikelihood
+    L2_loss2
     Φ
 end
 
 function LogDensityProblems.logdensity(ltd::PDELogTargetDensity, θ)
     # for parameter estimation neccesarry to use multioutput case
-    return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) + priorlogpdf(ltd, θ) +
-           L2LossData(ltd, θ)
+    if ltd.L2_loss2 === nothing
+        return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) +
+               priorlogpdf(ltd, θ) + L2LossData(ltd, θ)
+    else
+        return ltd.full_loglikelihood(setparameters(ltd, θ), ltd.allstd) +
+               priorlogpdf(ltd, θ) + L2LossData(ltd, θ) + ltd.L2_loss2(setparameters(ltd, θ), ltd.phynewstd)
+    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 = SymbolicUtils.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 = [[Symbolics.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 = [Symbolics.substitute.(
+                            masked_colloc_equation, Ref(rev_Dict_differentials))
+                        for masked_colloc_equation in masked_colloc_equations]
+    # 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
+    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 data_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
 
 @views function setparameters(ltd::PDELogTargetDensity, θ)
@@ -180,8 +253,8 @@ end
 """
     ahmc_bayesian_pinn_pde(pde_system, discretization;
         draw_samples = 1000, bcstd = [0.01], l2std = [0.05], phystd = [0.05],
-        priorsNNw = (0.0, 2.0), param = [], nchains = 1, Kernel = HMC(0.1, 30),
-        Adaptorkwargs = (Adaptor = StanHMCAdaptor,
+        phynewstd = [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)
@@ -210,6 +283,7 @@ end
   each dependant variable of interest.
 * `phystd`: Vector of standard deviations of BPINN prediction against Chosen Underlying PDE
   equations.
+* `phynewstd`: Vector of standard deviations of new loss term.
 * `priorsNNw`: Tuple of (mean, std) for BPINN Network parameters. Weights and Biases of
   BPINN are Normal Distributions by default.
 * `param`: Vector of chosen PDE's parameter's Distributions in case of Inverse problems.
@@ -235,14 +309,53 @@ 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), param = [], nchains = 1, Kernel = HMC(0.1, 30),
-        Adaptorkwargs = (Adaptor = StanHMCAdaptor,
+        phynewstd = [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
+        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)
+        # 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)]
+
+        # 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),
+                                        data_colloc_loss_functions[i],
+                                        nothing;
+                                        train_sets_pde = colloc_train_sets[i],
+                                        train_sets_bc = nothing)[1]
+                                    for i in eachindex(data_colloc_loss_functions)]
+
+        function L2_loss2(θ, phynewstd)
+            # 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
+
+    # add overall functionality for BC dataset points (case of parametric BC) ?
     if ((dataset_bc isa Nothing) && (dataset_pde isa Nothing))
         dataset = nothing
     elseif dataset_bc isa Nothing
@@ -306,8 +419,8 @@ function ahmc_bayesian_pinn_pde(pde_system, discretization;
 
     # dimensions would be total no of params,initial_nnθ for Lux namedTuples
     ℓπ = PDELogTargetDensity(
-        nparameters, strategy, dataset, priors, [phystd, bcstd, l2std],
-        names, ninv, initial_nnθ, full_weighted_loglikelihood, Φ)
+        nparameters, strategy, dataset, priors, [phystd, bcstd, l2std], phynewstd,
+        names, ninv, initial_nnθ, full_weighted_loglikelihood, newloss, Φ)
 
     Adaptor, Metric, targetacceptancerate = Adaptorkwargs[:Adaptor],
     Adaptorkwargs[:Metric], Adaptorkwargs[:targetacceptancerate]
@@ -320,8 +433,13 @@ 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",
+                ℓπ.L2_loss2(setparameters(ℓπ, initial_θ),
+                    ℓπ.phynewstd))
+        end
     end
 
     # parallel sampling option
@@ -370,11 +488,16 @@ function ahmc_bayesian_pinn_pde(pde_system, discretization;
 
         if verbose
             @printf("Sampling Complete.\n")
-            @printf("Current Physics Log-likelihood : %g\n",
+            @printf("Final Physics Log-likelihood : %g\n",
                 ℓπ.full_loglikelihood(setparameters(ℓπ, samples[end]), ℓπ.allstd))
-            @printf("Current Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, samples[end]))
-            @printf("Current MSE against dataset Log-likelihood : %g\n",
+            @printf("Final Prior Log-likelihood : %g\n", priorlogpdf(ℓπ, samples[end]))
+            @printf("Final SSE against dataset Log-likelihood : %g\n",
                 L2LossData(ℓπ, samples[end]))
+            if !(newloss isa Nothing)
+                @printf("Final new loss : %g\n",
+                    ℓπ.L2_loss2(setparameters(ℓπ, samples[end]),
+                        ℓπ.phynewstd))
+            end
         end
 
         fullsolution = BPINNstats(mcmc_chain, samples, stats)
diff --git a/src/advancedHMC_MCMC.jl b/src/advancedHMC_MCMC.jl
index 380d284f5..f7f18e09b 100644
--- a/src/advancedHMC_MCMC.jl
+++ b/src/advancedHMC_MCMC.jl
@@ -6,6 +6,7 @@
     dataset <: Union{Vector{Nothing}, Vector{<:Vector{<:AbstractFloat}}}
     priors <: Vector{<:Distribution}
     phystd::Vector{Float64}
+    phynewstd::Vector{Float64}
     l2std::Vector{Float64}
     autodiff::Bool
     physdt::Float64
@@ -97,7 +98,7 @@ suggested extra loss function for ODE solver case
     for i in 1:length(ltd.prob.u0)
         physlogprob += logpdf(
             MvNormal(deri_physsol[i, :],
-                Diagonal(abs2.(T(ltd.phystd[i]) .* ones(T, length(nnsol[i, :]))))),
+                Diagonal(abs2.(T(ltd.phynewstd[i]) .* ones(T, length(nnsol[i, :]))))),
             nnsol[i, :]
         )
     end
@@ -263,7 +264,7 @@ end
 """
     ahmc_bayesian_pinn_ode(prob, chain; strategy = GridTraining, dataset = [nothing],
                            init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0f0,
-                           l2std = [0.05], phystd = [0.05], priorsNNw = (0.0, 2.0),
+                           l2std = [0.05], phystd = [0.05], phynewstd = [0.05], priorsNNw = (0.0, 2.0),
                            param = [], nchains = 1, autodiff = false, Kernel = HMC,
                            Adaptorkwargs = (Adaptor = StanHMCAdaptor,
                                Metric = DiagEuclideanMetric, targetacceptancerate = 0.8),
@@ -336,6 +337,7 @@ Incase you are only solving the Equations for solution, do not provide dataset
   ~2/3 of draw samples)
 * `l2std`: standard deviation of BPINN prediction against L2 losses/Dataset
 * `phystd`: standard deviation of BPINN prediction against Chosen Underlying ODE System
+* `phynewstd`: standard deviation of new loss func term
 * `priorsNNw`: Tuple of (mean, std) for BPINN Network parameters. Weights and Biases of
   BPINN are Normal Distributions by default.
 * `param`: Vector of chosen ODE parameters Distributions in case of Inverse problems.
@@ -366,10 +368,10 @@ Incase you are only solving the Equations for solution, do not provide dataset
 function ahmc_bayesian_pinn_ode(
         prob::SciMLBase.ODEProblem, chain; strategy = GridTraining, dataset = [nothing],
         init_params = nothing, draw_samples = 1000, physdt = 1 / 20.0, l2std = [0.05],
-        phystd = [0.05], priorsNNw = (0.0, 2.0), param = [], nchains = 1, autodiff = false,
-        Kernel = HMC,
-        Adaptorkwargs = (Adaptor = StanHMCAdaptor, Metric = DiagEuclideanMetric,
-            targetacceptancerate = 0.8),
+        phystd = [0.05], phynewstd = [0.05], priorsNNw = (0.0, 2.0), param = [], nchains = 1,
+        autodiff = false, Kernel = HMC,
+        Adaptorkwargs = (Adaptor = StanHMCAdaptor,
+            Metric = DiagEuclideanMetric, targetacceptancerate = 0.8),
         Integratorkwargs = (Integrator = Leapfrog,), MCMCkwargs = (n_leapfrog = 30,),
         progress = false, verbose = false, estim_collocate = false)
     @assert !isinplace(prob) "The BPINN ODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t)."
@@ -415,16 +417,15 @@ function ahmc_bayesian_pinn_ode(
         nparameters += ninv
     end
 
-    t0 = prob.tspan[1]
     smodel = StatefulLuxLayer{true}(chain, nothing, st)
     # dimensions would be total no of params,initial_nnθ for Lux namedTuples
     ℓπ = LogTargetDensity(nparameters, prob, smodel, strategy, dataset, priors,
-        phystd, l2std, autodiff, physdt, ninv, initial_nnθ, estim_collocate)
+        phystd, phynewstd, l2std, autodiff, physdt, ninv, initial_nnθ, estim_collocate)
 
     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",
@@ -483,13 +484,13 @@ function ahmc_bayesian_pinn_ode(
 
         if verbose
             println("Sampling Complete.")
-            @printf("Current Physics Log-likelihood: %g\n",
+            @printf("Final Physics Log-likelihood: %g\n",
                 physloglikelihood(ℓπ, samples[end]))
-            @printf("Current Prior Log-likelihood: %g\n", priorweights(ℓπ, samples[end]))
-            @printf("Current MSE against dataset Log-likelihood: %g\n",
+            @printf("Final Prior Log-likelihood: %g\n", priorweights(ℓπ, samples[end]))
+            @printf("Final SSE against dataset Log-likelihood: %g\n",
                 L2LossData(ℓπ, samples[end]))
             if estim_collocate
-                @printf("Current gradient loss against dataset Log-likelihood: %g\n",
+                @printf("Final gradient loss against dataset Log-likelihood: %g\n",
                     L2loss2(ℓπ, samples[end]))
             end
         end
diff --git a/src/discretize.jl b/src/discretize.jl
index bed027aa2..43653ba7c 100644
--- a/src/discretize.jl
+++ b/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
@@ -538,23 +542,26 @@ function SciMLBase.symbolic_discretize(pde_system::PDESystem, discretization::Ab
             nothing, nothing
         end
 
+        # this includes losses from dataset domain points as well as discretization points
         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 += [logpdf(Normal(0, stdpdes[j]), pde_loss_function(θ))
-                                       for (j, pde_loss_function) in enumerate(datapde_loss_functions)]
+                pde_loglikelihoods += sum([pde_loglike_function(θ, stdpdes[j])
+                                           for (j, pde_loglike_function) in enumerate(datapde_loss_functions)])
             end
 
             if !(databc_loss_functions isa Nothing)
-                bc_loglikelihoods += [logpdf(Normal(0, stdbcs[j]), bc_loss_function(θ))
-                                      for (j, bc_loss_function) in enumerate(databc_loss_functions)]
+                bc_loglikelihoods += sum([bc_loglike_function(θ, stdbcs[j])
+                                          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
diff --git a/src/training_strategies.jl b/src/training_strategies.jl
index ca07676f2..dc466f091 100644
--- a/src/training_strategies.jl
+++ b/src/training_strategies.jl
@@ -14,26 +14,81 @@ corresponding to the grid spacing in each dimension.
     dx
 end
 
-# include dataset points in pde_residual loglikelihood (BayesianPINN)
+# dataset must have depvar values for same values of indvars
+function get_dataset_train_points(eqs, train_sets, pinnrep)
+    dict_depvar_input = pinnrep.dict_depvar_input
+    depvars = pinnrep.depvars
+    dict_depvars = pinnrep.dict_depvars
+    dict_indvars = pinnrep.dict_indvars
+
+    symbols_input = [(i, dict_depvar_input[i]) for i in depvars]
+    # [(:u, [:t])]
+    eq_args = NeuralPDE.get_argument(eqs, dict_indvars, dict_depvars)
+    # equation wise indvar presence ~ [[:t]]
+    # in each equation atleast one depvars must be a function of all indvars(to cover heterogenous/not case)
+
+    # train_sets follows order of depvars
+    # take dataset indvar values if for equations depvar's indvar matches input symbol indvar
+    points = []
+    for eq_arg in eq_args
+        eq_points = []
+        for i in eachindex(symbols_input)
+            if symbols_input[i][2] == eq_arg
+                push!(eq_points, train_sets[i][:, 2:end]')
+                # Terminate to avoid repetitive ind var points inclusion
+                break
+            end
+        end
+        # Concatenate points for this equation argument
+        push!(points, vcat(eq_points...))
+    end
+
+    return points
+end
+
+# includes dataset points in pde_residual loglikelihood (only for BayesianPINN)
 function merge_strategy_with_loglikelihood_function(pinnrep::PINNRepresentation,
         strategy::GridTraining, datafree_pde_loss_function,
         datafree_bc_loss_function; train_sets_pde = nothing, train_sets_bc = nothing)
     eltypeθ = recursive_eltype(pinnrep.flat_init_params)
     adaptor = EltypeAdaptor{eltypeθ}()
 
-    # 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)
+    # 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 = 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
+
     pde_loss_functions = if train_sets_pde !== nothing
-        pde_train_sets = [train_set[:, 2:end] for train_set in train_sets_pde] |> adaptor
-        [get_loss_function(pinnrep, _loss, _set, eltypeθ, strategy)
+        # 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
-        [get_loss_function(pinnrep, _loss, _set, eltypeθ, strategy)
+        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
         nothing
@@ -42,6 +97,19 @@ function merge_strategy_with_loglikelihood_function(pinnrep::PINNRepresentation,
     return pde_loss_functions, bc_loss_functions
 end
 
+function get_points_loss_functions(loss_function, train_set, eltypeθ, strategy::GridTraining;
+        τ = nothing)
+        # loss_function length is number of all points loss is being evaluated upon
+        # train sets rows are for each indvar, cols are coordinates (row_1,row_2,..row_n) at which loss evaluated
+    function loss(θ, std)
+        logpdf(
+            MvNormal(loss_function(train_set, θ)[1, :],
+                Diagonal(abs2.(std .* ones(size(train_set)[2])))),
+            zeros(size(train_set)[2]))
+    end
+end
+
+# only for PhysicsInformedNN
 function merge_strategy_with_loss_function(pinnrep::PINNRepresentation,
         strategy::GridTraining, datafree_pde_loss_function, datafree_bc_loss_function)
     (; domains, eqs, bcs, dict_indvars, dict_depvars) = pinnrep
diff --git a/test/BPINN_PDE_tests.jl b/test/BPINN_PDE_tests.jl
index 6a768533d..f5f5a96f7 100644
--- a/test/BPINN_PDE_tests.jl
+++ b/test/BPINN_PDE_tests.jl
@@ -1,4 +1,4 @@
-@testitem "BPINN PDE I: 2D Periodic System" tags=[:pdebpinn] begin
+@testitem "BPINN PDE I: 1D Periodic System" tags=[:pdebpinn] begin
     using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
           AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
           ComponentArrays
@@ -21,7 +21,7 @@
     discretization = BayesianPINN([chainl], GridTraining([0.01]))
 
     sol1 = ahmc_bayesian_pinn_pde(
-        pde_system, discretization; draw_samples = 1500, bcstd = [0.02],
+        pde_system, discretization; draw_samples = 1500, bcstd = [0.01],
         phystd = [0.01], priorsNNw = (0.0, 1.0), saveats = [1 / 50.0])
 
     analytic_sol_func(u0, t) = u0 + sinpi(2t) / (2pi)
@@ -29,8 +29,8 @@
     u_real = [analytic_sol_func(0.0, t) for t in ts]
     u_predict = pmean(sol1.ensemblesol[1])
 
-    @test u_predict≈u_real atol=0.5
-    @test mean(u_predict .- u_real) < 0.1
+    # absol tests
+    @test mean(abs, u_predict .- u_real) < 5e-2
 end
 
 @testitem "BPINN PDE II: 1D ODE" tags=[:pdebpinn] begin
@@ -240,7 +240,7 @@ end
     bcs = [u(0) ~ 0.0]
     domains = [t ∈ Interval(0.0, 2.0)]
 
-    chainl = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 1))
+    chainl = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1))
     initl, st = Lux.setup(Random.default_rng(), chainl)
 
     @named pde_system = PDESystem(eqs,
@@ -257,9 +257,10 @@ end
     u = u .+ (u .* 0.2) .* randn(size(u))
     dataset = [hcat(u, timepoints)]
 
+    # BPINNs are formulated with a mesh that must stay the same throughout sampling (as of now)
     @testset "$(nameof(typeof(strategy)))" for strategy in [
-        StochasticTraining(200),
-        QuasiRandomTraining(200),
+    # StochasticTraining(200),
+    # QuasiRandomTraining(200),
         GridTraining([0.02])
     ]
         discretization = BayesianPINN([chainl], strategy; param_estim = true,
@@ -268,8 +269,8 @@ end
         sol1 = ahmc_bayesian_pinn_pde(pde_system,
             discretization;
             draw_samples = 1500,
-            bcstd = [0.05],
-            phystd = [0.01], l2std = [0.01],
+            bcstd = [0.02],
+            phystd = [0.02], l2std = [0.02],
             priorsNNw = (0.0, 1.0),
             saveats = [1 / 50.0],
             param = [LogNormal(6.0, 0.5)])
@@ -279,9 +280,8 @@ end
         u_real = [analytic_sol_func1(0.0, t) for t in ts]
         u_predict = pmean(sol1.ensemblesol[1])
 
-        @test u_predict≈u_real atol=1.5
-        @test mean(u_predict .- u_real) < 0.1
-        @test sol1.estimated_de_params[1]≈param atol=param * 0.3
+        @test mean(abs, u_predict .- u_real) < 5e-2
+        @test sol1.estimated_de_params[1]≈param rtol=0.1
     end
 end
 
@@ -327,7 +327,7 @@ end
     ts = sol.t
     us = hcat(sol.u...)
     us = us .+ ((0.05 .* randn(size(us))) .* us)
-    ts_ = hcat(sol(ts).t...)[1, :]
+    ts_ = hcat(ts...)[1, :]
     dataset = [hcat(us[i, :], ts_) for i in 1:3]
 
     discretization = BayesianPINN(chain, GridTraining([0.01]); param_estim = true,
@@ -351,3 +351,179 @@ end
     @test sum(abs, pmean(p_) - 10.00) < 0.3 * idealp[1]
     # @test sum(abs, pmean(p_[2]) - (8 / 3)) < 0.3 * idealp[2]
 end
+
+@testitem "BPINN PDE Inv III: Improved Parametric Kuromo-Sivashinsky Equation solve" tags=[:pdebpinn] begin
+    using MCMCChains, Lux, ModelingToolkit, Distributions, OrdinaryDiffEq,
+          AdvancedHMC, Statistics, Random, Functors, NeuralPDE, MonteCarloMeasurements,
+          ComponentArrays
+    import ModelingToolkit: Interval, infimum, supremum
+
+    Random.seed!(100)
+
+    function recur_expression(exp, Dict_differentials)
+        for in_exp in exp.args
+            if !(in_exp isa Expr)
+                # skip +,== symbols, characters etc
+                continue
+
+            elseif in_exp.args[1] isa ModelingToolkit.Differential
+                # first symbol of differential term
+                # Dict_differentials for masking differential terms
+                # and resubstituting differentials in equations after putting in interpolations
+                # temp = in_exp.args[end]
+                Dict_differentials[eval(in_exp)] = Symbolics.variable("diff_$(length(Dict_differentials) + 1)")
+                return
+            else
+                recur_expression(in_exp, Dict_differentials)
+            end
+        end
+    end
+
+    @parameters x, t, α
+    @variables u(..)
+    Dt = Differential(t)
+    Dx = Differential(x)
+    Dx2 = Differential(x)^2
+    Dx3 = Differential(x)^3
+    Dx4 = Differential(x)^4
+
+    # α = 1 (KS equation to be parametric in a)
+    β = 4
+    γ = 1
+    eq = Dt(u(x, t)) + u(x, t) * Dx(u(x, t)) + α * Dx2(u(x, t)) + β * Dx3(u(x, t)) + γ * Dx4(u(x, t)) ~ 0
+
+    u_analytic(x, t; z = -x / 2 + t) = 11 + 15 * tanh(z) - 15 * tanh(z)^2 - 15 * tanh(z)^3
+    du(x, t; z = -x / 2 + t) = 15 / 2 * (tanh(z) + 1) * (3 * tanh(z) - 1) * sech(z)^2
+
+    bcs = [u(x, 0) ~ u_analytic(x, 0),
+        u(-10, t) ~ u_analytic(-10, t),
+        u(10, t) ~ u_analytic(10, t),
+        Dx(u(-10, t)) ~ du(-10, t),
+        Dx(u(10, t)) ~ du(10, t)]
+
+    # Space and time domains
+    domains = [x ∈ Interval(-10.0, 10.0),
+        t ∈ Interval(0.0, 1.0)]
+
+    # Discretization
+    dx = 0.4
+    dt = 0.2
+
+    # Function to compute analytical solution at a specific point (x, t)
+    function u_analytic_point(x, t)
+        z = -x / 2 + t
+        return 11 + 15 * tanh(z) - 15 * tanh(z)^2 - 15 * tanh(z)^3
+    end
+
+    # Function to generate the dataset matrix
+    function generate_dataset_matrix(domains, dx, dt, xlim, tlim)
+        x_values = xlim[1]:dx:xlim[2]
+        t_values = tlim[1]:dt:tlim[2]
+
+        dataset = []
+
+        for t in t_values
+            for x in x_values
+                u_value = u_analytic_point(x, t)
+                push!(dataset, [u_value, x, t])
+            end
+        end
+
+        return vcat([data' for data in dataset]...)
+    end
+
+    # considering sparse dataset from half of x's domain
+    datasetpde_new = [generate_dataset_matrix(domains, dx, dt, [-10, 0], [0.0, 1.0])]
+
+    # Adding Gaussian noise with a 0.8 std
+    noisydataset_new = deepcopy(datasetpde_new)
+    noisydataset_new[1][:, 1] = noisydataset_new[1][:, 1] .+
+                                (randn(size(noisydataset_new[1][:, 1])) .* 0.8)
+
+    # Neural network
+    chain = Lux.Chain(Lux.Dense(2, 8, Lux.tanh),
+        Lux.Dense(8, 8, Lux.tanh),
+        Lux.Dense(8, 1))
+
+    # Discretization for old and new models
+    discretization = NeuralPDE.BayesianPINN([chain],
+        GridTraining([dx, dt]), param_estim = true, dataset = [noisydataset_new, nothing])
+
+    # let α default to 2.0
+    @named pde_system = PDESystem(eq,
+        bcs,
+        domains,
+        [x, t],
+        [u(x, t)],
+        [α],
+        defaults = Dict([α => 2.0]))
+
+    # neccesarry for loss function construction (involves Operator masking) 
+    eqs = pde_system.eqs
+    Dict_differentials = Dict()
+    exps = toexpr.(eqs)
+    nullobj = [recur_expression(exp, Dict_differentials) for exp in exps]
+
+    # Dict_differentials is now ;
+    # Dict{Any, Any} with 5 entries:
+    #   Differential(x)(Differential(x)(u(x, t)))            => diff_5
+    #   Differential(x)(Differential(x)(Differential(x)(u(x… => diff_1
+    #   Differential(x)(Differential(x)(Differential(x)(Dif… => diff_2
+    #   Differential(x)(u(x, t))                             => diff_4
+    #   Differential(t)(u(x, t))                             => diff_3
+
+    # using HMC algorithm due to convergence, stability, time of training. (refer to mcmc chain plots)
+    # choice of std for objectives is very important
+    # pass in Dict_differentials, phystdnew arguments when using the new model
+
+    sol_new = ahmc_bayesian_pinn_pde(pde_system,
+        discretization;
+        draw_samples = 150,
+        bcstd = [0.1, 0.1, 0.1, 0.1, 0.1], phynewstd = [0.4],
+        phystd = [0.2], l2std = [0.8], param = [Distributions.Normal(2.0, 2)],
+        priorsNNw = (0.0, 1.0),
+        saveats = [1 / 100.0, 1 / 100.0],
+        Dict_differentials = Dict_differentials)
+
+    sol_old = ahmc_bayesian_pinn_pde(pde_system,
+        discretization;
+        draw_samples = 150,
+        bcstd = [0.1, 0.1, 0.1, 0.1, 0.1],
+        phystd = [0.2], l2std = [0.8], param = [Distributions.Normal(2.0, 2)],
+        priorsNNw = (0.0, 1.0),
+        saveats = [1 / 100.0, 1 / 100.0])
+
+    phi = discretization.phi[1]
+    xs, ts = [infimum(d.domain):dx:supremum(d.domain)
+              for (d, dx) in zip(domains, [dx / 10, dt])]
+    u_real = [[u_analytic(x, t) for x in xs] for t in ts]
+
+    u_predict_new = [[first(pmean(phi([x, t], sol_new.estimated_nn_params[1]))) for x in xs]
+                     for t in ts]
+
+    diff_u_new = [[abs(u_analytic(x, t) -
+                       first(pmean(phi([x, t], sol_new.estimated_nn_params[1]))))
+                   for x in xs]
+                  for t in ts]
+
+    u_predict_old = [[first(pmean(phi([x, t], sol_old.estimated_nn_params[1]))) for x in xs]
+                     for t in ts]
+    diff_u_old = [[abs(u_analytic(x, t) -
+                       first(pmean(phi([x, t], sol_old.estimated_nn_params[1]))))
+                   for x in xs]
+                  for t in ts]
+
+    unsafe_comparisons(true)
+    @test all(all, [((diff_u_new[i]) .^ 2 .< 0.8) for i in 1:6]) == true
+    @test all(all, [((diff_u_old[i]) .^ 2 .< 0.8) for i in 1:6]) == false
+
+    MSE_new = [mean(abs2, diff_u_new[i]) for i in 1:6]
+    MSE_old = [mean(abs2, diff_u_old[i]) for i in 1:6]
+    @test (MSE_new .< MSE_old) == [1, 1, 1, 1, 1, 1]
+
+    param_new = sol_new.estimated_de_params[1]
+    param_old = sol_old.estimated_de_params[1]
+    α = 1
+    @test abs(param_new - α) < 0.2 * α
+    @test abs(param_new - α) < abs(param_old - α)
+end
\ No newline at end of file
diff --git a/test/BPINN_tests.jl b/test/BPINN_tests.jl
index 7f1df5691..f2a33355a 100644
--- a/test/BPINN_tests.jl
+++ b/test/BPINN_tests.jl
@@ -137,7 +137,7 @@ end
     sol = solve(prob, Tsit5(); saveat = 0.1)
     u = sol.u
     time = sol.t
-    x̂ = u .+ (u .* 0.2) .* randn(size(u))
+    x̂ = u .+ (u .* 0.1) .* randn(size(u))
     dataset = [x̂, time]
     physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)]
 
@@ -148,16 +148,16 @@ end
     chainlux12 = Chain(Dense(1, 6, tanh), Dense(6, 6, tanh), Dense(6, 1))
     θinit, st = Lux.setup(Random.default_rng(), chainlux12)
 
+    # this a forward solve
     fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode(
-        prob, chainlux12, draw_samples = 1500, l2std = [0.03],
-        phystd = [0.03], priorsNNw = (0.0, 10.0))
+        prob, chainlux12, draw_samples = 500, phystd = [0.01], priorsNNw = (0.0, 10.0))
 
     fh_mcmc_chainlux22, fhsampleslux22, fhstatslux22 = ahmc_bayesian_pinn_ode(
-        prob, chainlux12, dataset = dataset, draw_samples = 1500, l2std = [0.03],
-        phystd = [0.03], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)])
+        prob, chainlux12, dataset = dataset, draw_samples = 500, l2std = [0.02],
+        phystd = [0.05], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)])
 
-    alg = BNNODE(chainlux12, dataset = dataset, draw_samples = 1500, l2std = [0.03],
-        phystd = [0.03], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)])
+    alg = BNNODE(chainlux12, dataset = dataset, draw_samples = 500, l2std = [0.02],
+        phystd = [0.05], priorsNNw = (0.0, 10.0), param = [Normal(-7, 4)])
 
     sol3lux_pestim = solve(prob, alg)
 
@@ -166,32 +166,25 @@ end
     #------------------------------ ahmc_bayesian_pinn_ode() call
     # Mean of last 500 sampled parameter's curves(lux chains)[Ensemble predictions]
     θ = [vector_to_parameters(fhsampleslux12[i], θinit)
-         for i in 1000:length(fhsampleslux12)]
+         for i in 400:length(fhsampleslux12)]
     luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)]
     luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)]
     meanscurve2_1 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean
 
     θ = [vector_to_parameters(fhsampleslux22[i][1:(end - 1)], θinit)
-         for i in 1000:length(fhsampleslux22)]
+         for i in 400:length(fhsampleslux22)]
     luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)]
     luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)]
     meanscurve2_2 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean
 
-    @test mean(abs, sol.u .- meanscurve2_1) < 1e-1
-    @test mean(abs, physsol1 .- meanscurve2_1) < 1e-1
-    @test mean(abs, sol.u .- meanscurve2_2) < 5e-2
-    @test mean(abs, physsol1 .- meanscurve2_2) < 5e-2
+    @test mean(abs, sol.u .- meanscurve2_1) < 1e-2
+    @test mean(abs, physsol1 .- meanscurve2_1) < 1e-2
+    @test mean(abs, sol.u .- meanscurve2_2) < 1.5
+    @test mean(abs, physsol1 .- meanscurve2_2) < 1.5
 
     # estimated parameters(lux chain)
-    param1 = mean(i[62] for i in fhsampleslux22[1000:length(fhsampleslux22)])
-    @test abs(param1 - p) < abs(0.3 * p)
-
-    #-------------------------- solve() call
-    # (lux chain)
-    @test mean(abs, physsol2 .- pmean(sol3lux_pestim.ensemblesol[1])) < 0.15
-    # estimated parameters(lux chain)
-    param1 = sol3lux_pestim.estimated_de_params[1]
-    @test abs(param1 - p) < abs(0.45 * p)
+    param1 = mean(i[62] for i in fhsampleslux22[400:length(fhsampleslux22)])
+    @test abs(param1 - p) < abs(0.5 * p)
 end
 
 @testitem "BPINN ODE: Translating from Flux" tags=[:odebpinn] begin
@@ -246,67 +239,117 @@ end
     sol = solve(prob, Tsit5(); saveat = 0.1)
     u = sol.u
     time = sol.t
-    x̂ = u .+ (0.3 .* randn(size(u)))
+    x̂ = u .+ (0.1 .* randn(size(u)))
     dataset = [x̂, time]
     physsol1 = [linear_analytic(prob.u0, p, time[i]) for i in eachindex(time)]
 
-    # separate set of points for testing the solve() call (it uses saveat 1/50 hence here length 501)
-    time1 = vec(collect(Float64, range(tspan[1], tspan[2], length = 501)))
-    physsol2 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)]
-
-    chainlux12 = Chain(Dense(1, 6, tanh), Dense(6, 6, tanh), Dense(6, 1))
+    chainlux12 = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1))
     θinit, st = Lux.setup(Random.default_rng(), chainlux12)
 
-    fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode(
-        prob, chainlux12, dataset = dataset, draw_samples = 1000, l2std = [0.1],
-        phystd = [0.03], priorsNNw = (0.0, 1.0), param = [Normal(-7, 3)])
-
     fh_mcmc_chainlux22, fhsampleslux22, fhstatslux22 = ahmc_bayesian_pinn_ode(
-        prob, chainlux12, dataset = dataset, draw_samples = 1000,
-        l2std = [0.1], phystd = [0.03], priorsNNw = (0.0, 1.0),
-        param = [Normal(-7, 3)], estim_collocate = true)
+        prob, chainlux12,
+        dataset = dataset,
+        draw_samples = 500,
+        l2std = [0.1],
+        phystd = [0.01],
+        phynewstd = [0.01],
+        priorsNNw = (0.0,
+            1.0),
+        param = [
+            Normal(-7, 3)
+        ], estim_collocate = true)
 
-    alg = BNNODE(
-        chainlux12, dataset = dataset, draw_samples = 1000, l2std = [0.1], phystd = [0.03],
-        priorsNNw = (0.0, 1.0), param = [Normal(-7, 3)], estim_collocate = true)
-
-    sol3lux_pestim = solve(prob, alg)
+    fh_mcmc_chainlux12, fhsampleslux12, fhstatslux12 = ahmc_bayesian_pinn_ode(
+        prob, chainlux12,
+        dataset = dataset,
+        draw_samples = 500,
+        l2std = [0.1],
+        phystd = [0.01],
+        priorsNNw = (0.0,
+            1.0),
+        param = [
+            Normal(-7, 3)
+        ])
 
     # testing timepoints
     t = sol.t
     #------------------------------ ahmc_bayesian_pinn_ode() call
-    # Mean of last 500 sampled parameter's curves(lux chains)[Ensemble predictions]
+    # Mean of last 100 sampled parameter's curves(lux chains)[Ensemble predictions]
     θ = [vector_to_parameters(fhsampleslux12[i][1:(end - 1)], θinit)
-         for i in 750:length(fhsampleslux12)]
+         for i in 400:length(fhsampleslux12)]
     luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)]
     luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)]
     meanscurve2_1 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean
 
     θ = [vector_to_parameters(fhsampleslux22[i][1:(end - 1)], θinit)
-         for i in 750:length(fhsampleslux22)]
+         for i in 400:length(fhsampleslux22)]
     luxar = [chainlux12(t', θ[i], st)[1] for i in eachindex(θ)]
     luxmean = [mean(vcat(luxar...)[:, i]) for i in eachindex(t)]
     meanscurve2_2 = prob.u0 .+ (t .- prob.tspan[1]) .* luxmean
 
-    @test_broken mean(abs.(sol.u .- meanscurve2_2)) < 6e-2
-    @test_broken mean(abs.(physsol1 .- meanscurve2_2)) < 6e-2
+    @test mean(abs.(sol.u .- meanscurve2_2)) < 1e-2
+    @test mean(abs.(physsol1 .- meanscurve2_2)) < 1e-2
     @test mean(abs.(sol.u .- meanscurve2_1)) > mean(abs.(sol.u .- meanscurve2_2))
     @test mean(abs.(physsol1 .- meanscurve2_1)) > mean(abs.(physsol1 .- meanscurve2_2))
 
     # estimated parameters(lux chain)
-    param2 = mean(i[62] for i in fhsampleslux22[750:length(fhsampleslux22)])
-    @test_broken abs(param2 - p) < abs(0.25 * p)
+    param2 = mean(i[62] for i in fhsampleslux22[400:length(fhsampleslux22)])
+    @test abs(param2 - p) < abs(0.05 * p)
 
-    param1 = mean(i[62] for i in fhsampleslux12[750:length(fhsampleslux12)])
-    @test abs(param1 - p) < abs(0.8 * p)
+    param1 = mean(i[62] for i in fhsampleslux12[400:length(fhsampleslux12)])
+    @test abs(param1 - p) > abs(0.5 * p)
     @test abs(param2 - p) < abs(param1 - p)
+end
+
+@testitem "BPINN ODE III: new objective solve call" tags=[:odebpinn] begin
+    using MCMCChains, Distributions, OrdinaryDiffEq, OptimizationOptimisers, Lux,
+          AdvancedHMC, Statistics, Random, Functors, ComponentArrays, MonteCarloMeasurements
+    import Flux
+
+    Random.seed!(100)
+
+    linear = (u, p, t) -> u / p + exp(t / p) * cos(t)
+    tspan = (0.0, 10.0)
+    u0 = 0.0
+    p = -5.0
+    prob = ODEProblem(linear, u0, tspan, p)
+    linear_analytic = (u0, p, t) -> exp(t / p) * (u0 + sin(t))
+
+    # SOLUTION AND CREATE DATASET
+    sol = solve(prob, Tsit5(); saveat = 0.1)
+    u = sol.u
+    time = sol.t
+    x̂ = u .+ (0.1 .* randn(size(u)))
+    dataset = [x̂, time]
+
+    # set of points for testing the solve() call (it uses saveat 1/50 hence here length 501)
+    time1 = vec(collect(Float64, range(tspan[1], tspan[2], length = 501)))
+    physsol2 = [linear_analytic(prob.u0, p, time1[i]) for i in eachindex(time1)]
+
+    chainlux12 = Lux.Chain(Lux.Dense(1, 6, tanh), Lux.Dense(6, 6, tanh), Lux.Dense(6, 1))
+    θinit, st = Lux.setup(Random.default_rng(), chainlux12)
+
+    alg = BNNODE(chainlux12,
+        dataset = dataset,
+        draw_samples = 1000,
+        l2std = [0.1],
+        phystd = [0.01],
+        phynewstd = [0.01],
+        priorsNNw = (0.0,
+            1.0),
+        param = [
+            Normal(-7, 3)
+        ], numensemble = 200,
+        estim_collocate = true)
+
+    sol3lux_pestim = solve(prob, alg)
 
     #-------------------------- solve() call
-    # (lux chain)
-    @test_broken mean(abs.(physsol2 .- pmean(sol3lux_pestim.ensemblesol[1]))) < 0.1
-    # estimated parameters(lux chain)
+    @test mean(abs.(physsol2 .- pmean(sol3lux_pestim.ensemblesol[1]))) < 1e-2
+
+    # estimated parameters
     param3 = sol3lux_pestim.estimated_de_params[1]
-    @test_broken abs(param3 - p) < abs(0.2 * p)
+    @test abs(param3 - p) < abs(0.05 * p)
 end
 
 @testitem "BPINN ODE IV: Improvement" tags=[:odebpinn] begin
@@ -318,43 +361,56 @@ end
 
     function lotka_volterra(u, p, t)
         # Model parameters.
-        α, β, γ, δ = p
+        α, δ = p
         # Current state.
         x, y = u
 
         # Evaluate differential equations.
-        dx = (α - β * y) * x # prey
-        dy = (δ * x - γ) * y # predator
+        dx = (α - y) * x # prey
+        dy = (x - δ) * y  # predator
 
         return [dx, dy]
     end
 
     # initial-value problem.
     u0 = [1.0, 1.0]
-    p = [1.5, 1.0, 3.0, 1.0]
-    tspan = (0.0, 4.0)
+    p = [1.5, 3.0]
+    tspan = (0.0, 7.0)
     prob = ODEProblem(lotka_volterra, u0, tspan, p)
 
-    # Solve using OrdinaryDiffEq.jl solver
-    dt = 0.2
+    # OrdinaryDiffEq.jl solve
+    dt = 0.1
     solution = solve(prob, Tsit5(); saveat = dt)
 
     times = solution.t
     u = hcat(solution.u...)
-    x = u[1, :] + (0.8 .* randn(length(u[1, :])))
-    y = u[2, :] + (0.8 .* randn(length(u[2, :])))
+    x = u[1, :] + (0.5 .* randn(length(u[1, :])))
+    y = u[2, :] + (0.5 .* randn(length(u[2, :])))
     dataset = [x, y, times]
 
-    chain = Chain(Dense(1, 6, tanh), Dense(6, 6, tanh), Dense(6, 2))
-
-    alg1 = BNNODE(chain; dataset = dataset, draw_samples = 1000,
-        l2std = [0.2, 0.2], phystd = [0.1, 0.1], priorsNNw = (0.0, 1.0),
-        param = [Normal(2, 0.5), Normal(2, 0.5), Normal(2, 0.5), Normal(2, 0.5)])
-
-    alg2 = BNNODE(chain; dataset = dataset, draw_samples = 1000,
-        l2std = [0.2, 0.2], phystd = [0.1, 0.1], priorsNNw = (0.0, 1.0),
-        param = [Normal(2, 0.5), Normal(2, 0.5), Normal(2, 0.5), Normal(2, 0.5)],
-        estim_collocate = true)
+    chain = Lux.Chain(Lux.Dense(1, 7, tanh), Lux.Dense(7, 7, tanh),
+        Lux.Dense(7, 2))
+
+    alg1 = BNNODE(chain;
+        dataset = dataset,
+        draw_samples = 1000,
+        l2std = [0.5, 0.5],
+        phystd = [0.5, 0.5],
+        priorsNNw = (0.0, 1.0),
+        param = [
+            Normal(2, 2),
+            Normal(2, 2)])
+
+    alg2 = BNNODE(chain;
+        dataset = dataset,
+        draw_samples = 1000,
+        l2std = [0.5, 0.5],
+        phystd = [0.5, 0.5],
+        phynewstd = [1.0, 1.0],
+        priorsNNw = (0.0, 1.0),
+        param = [
+            Normal(2, 2),
+            Normal(2, 2)], estim_collocate = true)
 
     @time sol_pestim1 = solve(prob, alg1; saveat = dt)
     @time sol_pestim2 = solve(prob, alg2; saveat = dt)
@@ -362,5 +418,12 @@ end
     unsafe_comparisons(true)
     bitvec = abs.(p .- sol_pestim1.estimated_de_params) .>
              abs.(p .- sol_pestim2.estimated_de_params)
-    @test_broken bitvec == ones(size(bitvec))
-end
+    @test bitvec == ones(size(bitvec))
+
+    Loss_1 = mean(abs, u[1, :] .- pmean(sol_pestim1.ensemblesol[1])) +
+             mean(abs, u[2, :] .- pmean(sol_pestim1.ensemblesol[2]))
+    Loss_2 = mean(abs, u[1, :] .- pmean(sol_pestim2.ensemblesol[1])) +
+             mean(abs, u[2, :] .- pmean(sol_pestim2.ensemblesol[2]))
+
+    @test Loss_1 > Loss_2
+end
\ No newline at end of file