make PoiintSoilver use Problems logdensity function

Author: bgctw

dev/doubleMM.jl

@@ -137,7 +137,7 @@ end
     () -> begin # optimized loss is indeed lower than with true parameters
         int_ϕθP = ComponentArrayInterpreter(CA.ComponentVector(
             ϕg = 1:length(prob0.ϕg), θP = prob0.θP))
-        loss_gf = get_loss_gf(prob0.g, prob0.transM, prob0.transP, prob0.f, Float32[], int_ϕθP)
+        loss_gf = get_loss_gf(prob0.g, prob0.transM, prob0.transP, prob0.f, Float32[], py, int_ϕθP)
         loss_gf(vcat(prob3.ϕg, prob3.θP), xM, xP, y_o, y_unc, i_sites)[1]
         loss_gf(vcat(prob3o.ϕg, prob3o.θP), xM, xP, y_o, y_unc, i_sites)[1]
         #

src/HybridSolver.jl

@@ -32,13 +32,14 @@ function CommonSolve.solve(prob::AbstractHybridProblem, solver::HybridPointSolve
         train_loader_dev = train_loader
     end
     f = get_hybridproblem_PBmodel(prob; scenario, use_all_sites=false)
+    py = get_hybridproblem_neg_logden_obs(prob; scenario)
     pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario)
     n_site, n_batch = get_hybridproblem_n_site_and_batch(prob; scenario)
     priors = get_hybridproblem_priors(prob; scenario)
     priorsP = [priors[k] for k in keys(par_templates.θP)]
     priorsM = [priors[k] for k in keys(par_templates.θM)]
     #intP = ComponentArrayInterpreter(par_templates.θP)
-    loss_gf = get_loss_gf(g_dev, transM, transP, f,  intϕ;
+    loss_gf = get_loss_gf(g_dev, transM, transP, f,  py, intϕ;
         cdev=infer_cdev(gdevs), pbm_covars, n_site_batch=n_batch, priorsP, priorsM,)
     # call loss function once
     l1 = is_infer ? 

src/gf.jl

@@ -137,6 +137,7 @@ Create a loss function for given
 - g(x, ϕ): machine learning model 
 - transM: transforamtion of parameters at unconstrained space
 - f(θMs, θP): mechanistic model 
+- py: `function(y_pred, y_obs, y_unc)` to compute negative log-likelihood, i.e. cost
 - intϕ: interpreter attaching axis with components ϕg and ϕP
 - intP: interpreter attaching axis to ζP = ϕP with components used by f,
   The default, uses `intϕ(ϕ)` as a template
@@ -160,7 +161,7 @@ and returns a NamedTuple of
 - `neg_log_prior`: negative log-prior of `θMs` and `θP`
 - `neg_log_prior`: negative log-prior of `θMs` and `θP`
 """
-function get_loss_gf(g, transM, transP, f,  
+function get_loss_gf(g, transM, transP, f, py,  
     intϕ::AbstractComponentArrayInterpreter,
     intP::AbstractComponentArrayInterpreter = ComponentArrayInterpreter(
         intϕ(1:length(intϕ)).ϕP);
@@ -178,7 +179,6 @@ function get_loss_gf(g, transM, transP, f,
         #, intP = get_concrete(intP)
         #inv_transP = inverse(transP), kwargs = kwargs
         function loss_gf(ϕ, xM, xP, y_o, y_unc, i_sites)
-            σ = exp.(y_unc ./ 2)
             ϕc = intϕ(ϕ)
             # μ_ζP = ϕc.ϕP
             # xMP = _append_each_covars(xM, CA.getdata(μ_ζP), pbm_covar_indices)
@@ -190,7 +190,9 @@ function get_loss_gf(g, transM, transP, f,
             y_pred, θMs_pred, θP_pred = gf(
                 g, transMs, transP, f, xM, xP, CA.getdata(ϕc.ϕg), CA.getdata(ϕc.ϕP), 
                 pbm_covar_indices; cdev, kwargs...)
-            nLy = sum(abs2, (y_pred .- y_o) ./ σ) 
+            #σ = exp.(y_unc ./ 2)
+            #nLy = sum(abs2, (y_pred .- y_o) ./ σ) 
+            nLy = py( y_pred, y_o, y_unc)
             # logpdf is not typestable for Distribution{Univariate, Continuous}
             logpdf_t = (prior, θ) -> logpdf(prior, θ)::eltype(θP_pred)
             logpdf_tv = (prior, θ::AbstractVector) -> begin

src/logden_normal.jl

@@ -1,7 +1,7 @@
 """
     neg_logden_indep_normal(obs, μ, logσ2s; σfac=1.0) 
 
-Compute the negative Log-density of `θM` for multiple independent normal distributions,
+Compute the negative Log-density of `obs` for multiple independent normal distributions,
 given estimated means `μ` and estimated log of variance parameters `logσ2s`.
 
 All the arguments should be vectors of the same length.

test/test_HybridProblem.jl

@@ -141,6 +141,7 @@ test_without_flux = (scenario) -> begin
         train_loader = get_hybridproblem_train_dataloader(prob; scenario)
         (xM, xP, y_o, y_unc, i_sites) = first(train_loader)
         f = get_hybridproblem_PBmodel(prob; scenario, use_all_sites = false)
+        py = get_hybridproblem_neg_logden_obs(prob; scenario)
         par_templates = get_hybridproblem_par_templates(prob; scenario)
         #f(par_templates.θP, hcat(par_templates.θM, par_templates.θM), xP[1:2])
         (; transM, transP) = get_hybridproblem_transforms(prob; scenario)
@@ -154,7 +155,7 @@ test_without_flux = (scenario) -> begin
         p = p0 = vcat(ϕg0, par_templates.θP .* convert(eltype(ϕg0), 0.8))  
 
         # Pass the site-data for the batches as separate vectors wrapped in a tuple
-        loss_gf = get_loss_gf(g, transM, transP, f,  intϕ; 
+        loss_gf = get_loss_gf(g, transM, transP, f, py, intϕ; 
             pbm_covars, n_site_batch = n_batch, priorsP, priorsM)
         (_xM, _xP, _y_o, _y_unc, _i_sites) = first(train_loader)
         l1 = loss_gf(p0, _xM, _xP, _y_o, _y_unc, _i_sites)

test/test_doubleMM.jl

@@ -199,6 +199,7 @@ end
     n_site, n_site_batch = get_hybridproblem_n_site_and_batch(prob; scenario)
     f = get_hybridproblem_PBmodel(prob; scenario, use_all_sites = false)
     f2 = get_hybridproblem_PBmodel(prob; scenario, use_all_sites = true)
+    py = get_hybridproblem_neg_logden_obs(prob; scenario)
     priors = get_hybridproblem_priors(prob; scenario)
     priorsP = [priors[k] for k in keys(par_templates.θP)]
     priorsM = [priors[k] for k in keys(par_templates.θM)]
@@ -218,9 +219,9 @@ end
     pbm_covars = get_hybridproblem_pbmpar_covars(prob; scenario)
 
     #loss_gf = get_loss_gf(g, transM, f,  intϕ; gdev = identity)
-    loss_gf = get_loss_gf(g, transM, transP, f,  intϕ;
+    loss_gf = get_loss_gf(g, transM, transP, f,  py, intϕ;
         pbm_covars, n_site_batch = n_batch, priorsP, priorsM)
-    loss_gf2 = get_loss_gf(g, transM, transP, f2,  intϕ;
+    loss_gf_site = get_loss_gf(g, transM, transP, f2, py, intϕ;
         pbm_covars, n_site_batch = n_site, priorsP, priorsM)
     nLjoint = @inferred first(loss_gf(p0, first(train_loader)...))
     (xM_batch, xP_batch, y_o_batch, y_unc_batch, i_sites_batch) = first(train_loader)
@@ -237,7 +238,7 @@ end
         #optprob, Adam(0.02), callback = callback_loss(100), maxiters = 5000);
         optprob, Adam(0.02), maxiters = 2000)
 
-    (;nLjoint_pen, y_pred, θMs_pred, θP_pred, nLy, neg_log_prior, loss_penalty) = loss_gf2(
+    (;nLjoint_pen, y_pred, θMs_pred, θP_pred, nLy, neg_log_prior, loss_penalty) = loss_gf_site(
         res.u, train_loader.data...)
     #(nLjoint,  y_pred, θMs_pred, θP, nLy, neg_log_prior, loss_penalty) = loss_gf(p0, xM, xP, y_o, y_unc);
     θMs_pred = CA.ComponentArray(θMs_pred, CA.getaxes(θMs_true'))

test/test_util_gpu.jl

@@ -15,11 +15,11 @@ gdev = gpu_device()
 if gdev isa MLDataDevices.CUDADevice
     @testset "ones_similar_x" begin
         B = CUDA.rand(Float32, 5, 2);    # GPU matrix
-        @test HVI.ones_similar_x(B, size(B,1)) isa CuArray
-        @test HVI.ones_similar_x(ComponentVector(b=B), size(B,1)) isa CuArray
-        @test HVI.ones_similar_x(B', size(B,1)) isa CuArray
-        @test HVI.ones_similar_x(@view(B[:,2]), size(B,1)) isa CuArray
-        @test HVI.ones_similar_x(ComponentVector(b=B)[:,1], size(B,1)) isa CuArray
+        @test HVI.ones_similar_x(B, size(B,1)) isa CUDA.CuArray
+        @test HVI.ones_similar_x(ComponentVector(b=B), size(B,1)) isa CUDA.CuArray
+        @test HVI.ones_similar_x(B', size(B,1)) isa CUDA.CuArray
+        @test HVI.ones_similar_x(@view(B[:,2]), size(B,1)) isa CUDA.CuArray
+        @test HVI.ones_similar_x(ComponentVector(b=B)[:,1], size(B,1)) isa CUDA.CuArray
     end
 end