add method create_nsite_applicator to PBMApplicator

Author: bgctw

src/ComponentArrayInterpreter.jl

@@ -171,6 +171,7 @@ function ComponentArrayInterpreter(
     ComponentArrayInterpreter(n_dims, CA.getaxes(ca), m_dims)
 end
 
+
 function ComponentArrayInterpreter(
     n_dims::NTuple{N,<:Integer}, axes::NTuple{A,<:CA.AbstractAxis},
     m_dims::NTuple{M,<:Integer}) where {N,A,M}

src/HybridVariationalInference.jl

@@ -55,6 +55,7 @@ include("ModelApplicator.jl")
 
 export AbstractPBMApplicator, NullPBMApplicator, PBMSiteApplicator, PBMPopulationApplicator
 export DirectPBMApplicator, PBMPopulationGlobalApplicator
+export create_nsite_applicator
 include("PBMApplicator.jl")
 
 # export AbstractGPUDataHandler, NullGPUDataHandler, get_default_GPUHandler

src/PBMApplicator.jl

@@ -28,6 +28,11 @@ function (app::AbstractPBMApplicator)(θP::AbstractArray, θMs::AbstractArray, x
     apply_model(app, θP, θMs, xP)
 end
 
+function create_nsite_applicator(app::AbstractPBMApplicator, n_site) 
+    copy(app)
+end
+
+
 """
     apply_model(app::AbstractPBMApplicator, θsP::AbstractVector, θsMs::AbstractMatrix, xP::AbstractMatrix) 
     apply_model(app::AbstractPBMApplicator, θsP::AbstractMatrix, θsMs::AbstractArray{ET,3}, xP) 
@@ -114,7 +119,7 @@ struct PBMSiteApplicator{F, IT, IXT, VFT} <: AbstractPBMApplicator
 end
 
 """
-    PBMSiteApplicator(fθ, n_batch; θP, θM, θFix, xPvec)
+    PBMSiteApplicator(fθ; θP, θM, θFix, xPvec)
 
 Construct AbstractPBMApplicator from process-based model `fθ` that computes predictions
 for a single site.
@@ -189,10 +194,10 @@ end
 @functor PBMPopulationApplicator (θFixm, rep_fac)
 
 """
-    PBMPopulationApplicator(fθpop, n_batch; θP, θM, θFix, xPvec)
+    PBMPopulationApplicator(fθpop, n_site; θP, θM, θFix, xPvec)
 
 Construct AbstractPBMApplicator from process-based model `fθ` that computes predictions
-across sites for a population of size `n_batch`.
+across sites for a population of size `n_site`.
 The applicator combines enclosed `θFix`, with provided `θMs` and `θP`
 to a `ComponentMatrix` with parameters with one row for each site, that
 can be column-indexed by Symbols.
@@ -203,29 +208,41 @@ can be column-indexed by Symbols.
     - `θc`: parameters: `ComponentMatrix` (n_site x n_par) with each row a parameter vector
     - `xPc`: observations: `ComponentMatrix` (n_obs x n_site) with each column 
     observationsfor one site
-- `n_batch`: number of indiduals, i.e. rows in `θMs`
+- `n_site`: number of indiduals, i.e. rows in `θMs`
 - `θP`: `ComponentVector` template of global process model parameters
 - `θM`: `ComponentVector` template of individual process model parameters
 - `θFix`: `ComponentVector` of actual fixed process model parameters
 - `xPvec`: `ComponentVector` template of model drivers for a single site
 """
-function PBMPopulationApplicator(fθpop, n_batch; 
+function PBMPopulationApplicator(fθpop, n_site; 
     θP::CA.ComponentVector, θM::CA.ComponentVector, θFix::CA.ComponentVector, 
     xPvec::CA.ComponentVector
     )
     intθvec = ComponentArrayInterpreter(flatten1(CA.ComponentVector(; θP, θM, θFix)))
     int_xP_vec = ComponentArrayInterpreter(xPvec)
-    isFix = repeat(axes(θFix, 1)', n_batch)
+    isFix = repeat(axes(θFix, 1)', n_site)
     #
-    intθ = get_concrete(ComponentArrayInterpreter((n_batch,), intθvec))
-    int_xP = get_concrete(ComponentArrayInterpreter(int_xP_vec, (n_batch,)))
-    #isP = repeat(axes(θP, 1)', n_batch)
+    intθ = get_concrete(ComponentArrayInterpreter((n_site,), intθvec))
+    int_xP = get_concrete(ComponentArrayInterpreter(int_xP_vec, (n_site,)))
+    #isP = repeat(axes(θP, 1)', n_site)
     # n_site = size(θMs, 1)
-    rep_fac = ones_similar_x(θP, n_batch) # to reshape into matrix, avoiding repeat
+    rep_fac = ones_similar_x(θP, n_site) # to reshape into matrix, avoiding repeat
     θFixm = CA.ComponentMatrix(θFix[isFix], (CA.FlatAxis(), CA.getaxes(θFix)[1]))
     PBMPopulationApplicator(fθpop, θFixm, rep_fac, intθ, int_xP)        
 end
 
+function create_nsite_applicator(app::PBMPopulationApplicator, n_site) 
+    θFix = app.θFixm[1,:]
+    isFix = repeat(axes(θFix, 1)', n_site)
+    θFixm = CA.ComponentMatrix(θFix[isFix], (CA.FlatAxis(), CA.getaxes(θFix)[1]))
+    #
+    intθ = get_concrete(ComponentArrayInterpreter((n_site,), (CA.getaxes(app.intθ)[2],),()))
+    int_xP = get_concrete(ComponentArrayInterpreter(
+        (), (CA.getaxes(app.int_xP)[1],), (n_site,)))
+    rep_fac = ones_similar_x(θFix, n_site) # to reshape into matrix, avoiding repeat
+    PBMPopulationApplicator(app.fθpop, θFixm, rep_fac, intθ, int_xP)        
+end
+
 function apply_model(app::PBMPopulationApplicator, θP::AbstractVector, θMs::AbstractMatrix, xP) 
     if (CA.getdata(θP) isa GPUArraysCore.AbstractGPUArray) && 
         (!(CA.getdata(app.θFixm) isa GPUArraysCore.AbstractGPUArray) || 
@@ -262,10 +279,10 @@ end
 
 
 """
-    PBMPopulationGlobalApplicator(fθpop, n_batch; θP, θM, θFix, xPvec)
+    PBMPopulationGlobalApplicator(fθpop, n_site; θP, θM, θFix, xPvec)
 
 Construct AbstractPBMApplicator from process-based model `fθ` that computes predictions
-across sites for a population of size `n_batch`.
+across sites for a population of size `n_site`.
 The applicator combines enclosed `θFix`, with provided `θMs` and `θP`
 to a `ComponentMatrix` with parameters with one row for each site, that
 can be column-indexed by Symbols.
@@ -278,24 +295,33 @@ can be column-indexed by Symbols.
     - `θgc`: parameters: `ComponentVector` (n_par_global) 
     - `xPc`: observations: `ComponentMatrix` (n_obs x n_site) with each column 
     observationsfor one site
-- `n_batch`: number of indiduals, i.e. rows in `θMs`
+- `n_site`: number of indiduals, i.e. rows in `θMs`
 - `θP`: `ComponentVector` template of global process model parameters
 - `θM`: `ComponentVector` template of individual process model parameters
 - `θFix`: `ComponentVector` of actual fixed process model parameters
 - `xPvec`: `ComponentVector` template of model drivers for a single site
 """
-function PBMPopulationGlobalApplicator(fθpop, n_batch; 
+function PBMPopulationGlobalApplicator(fθpop, n_site; 
     θP::CA.ComponentVector, θM::CA.ComponentVector, θFix::CA.ComponentVector, 
     xPvec::CA.ComponentVector
     )
-    intθvec = ComponentArrayInterpreter(flatten1(CA.ComponentVector(; θP, θM, θFix)))
+    #intθvec = ComponentArrayInterpreter(flatten1(CA.ComponentVector(; θP, θM, θFix)))
     int_xP_vec = ComponentArrayInterpreter(xPvec)
-    intθs = get_concrete(ComponentArrayInterpreter((n_batch,), θM))
+    intθs = get_concrete(ComponentArrayInterpreter((n_site,), θM))
     intθg = get_concrete(ComponentArrayInterpreter(vcat(θP, θFix)))
-    int_xP = get_concrete(ComponentArrayInterpreter(int_xP_vec, (n_batch,)))
+    int_xP = get_concrete(ComponentArrayInterpreter(int_xP_vec, (n_site,)))
     PBMPopulationGlobalApplicator(fθpop, θFix, intθs, intθg, int_xP)        
 end
 
+function create_nsite_applicator(app::PBMPopulationGlobalApplicator, n_site) 
+    @info("called PBMPopulationGlobalApplicator.create_nsite_applicator")
+    intθs = get_concrete(ComponentArrayInterpreter((n_site,), (CA.getaxes(app.intθs)[2],),()))
+    int_xP = get_concrete(ComponentArrayInterpreter(
+        (), (CA.getaxes(app.int_xP)[1],), (n_site,)))
+    PBMPopulationGlobalApplicator(app.fθpop, app.θFix, intθs, app.intθg, int_xP)        
+end
+
+
 function apply_model(app::PBMPopulationGlobalApplicator, θP::AbstractVector, θMs::AbstractMatrix, xP) 
     if (CA.getdata(θP) isa GPUArraysCore.AbstractGPUArray) && 
         (!(CA.getdata(app.θFix) isa GPUArraysCore.AbstractGPUArray) || 

test/test_PBMApplicator.jl

@@ -6,6 +6,49 @@ import Zygote
 
 using MLDataDevices, CUDA, cuDNN, GPUArraysCore
 
+f_pop = function(θsc, xPc)
+    local n_obs = size(xPc, 1)
+    is_valid = isfinite.(CA.getdata(xPc))
+    a1 = is_valid .* CA.getdata(θsc[:,:a1])' 
+    a2 = is_valid .* CA.getdata(θsc[:,:a2])' 
+    # b in θP has been expanded in PopulationApplicator
+    b = is_valid .* CA.getdata(θsc[:,:b])' 
+    y = a1 .+  log.(a2) .* abs2.(cos.(b .- 0.2)) .* xPc.^2
+end
+
+() -> begin
+    include("test/test_scratch.jl")
+end
+
+@testset "PBMPopulationApplicator" begin
+    n_obs = 3
+    n_site = 5
+    xPvec = CA.ComponentVector(s1 = 1.0:n_obs)
+    xPc = xPvec .* ones(n_site)' .+ abs2.(randn(n_obs, n_site) .* 0.1)
+    θP = CA.ComponentVector(b=3.0)
+    θM = CA.ComponentVector(a1=2.0,a2=1.0)
+    θFix  = CA.ComponentVector(c=1.5)
+    #
+    θMs = (ones(n_site) .* θM') .+ abs2.(randn(n_site, length(θM)) .* 0.1)
+    θs = hcat(ones(n_site) .* θP', θMs)
+    y_obs = f_pop(θs, xPc)
+    g = PBMPopulationApplicator(f_pop, n_site; θP, θM, θFix, xPvec)
+    ret = g(θP, θMs, xPc)
+    @test ret ≈ y_obs
+
+    gr = Zygote.gradient((θP, θMs) -> sum(g(θP, θMs, xPc)), CA.getdata(θP), CA.getdata(θMs))
+
+    xPc_NaN = copy(xPc); xPc_NaN[2:3,1] .= NaN
+    ret2 = g(θP, θMs, xPc_NaN)
+    gr = Zygote.gradient((θP, θMs) -> sum(g(θP, θMs, xPc_NaN)), CA.getdata(θP), CA.getdata(θMs))
+    @test all(isfinite.(gr[1])) 
+
+    n_site_m = 4
+    gm = create_nsite_applicator(g, n_site_m)
+    retm = gm(θP, θMs[1:n_site_m,:], xPc[:,1:n_site_m])
+    @test retm ≈ y_obs[:,1:n_site_m]
+end;
+
 f_global = function(θsc, θgc, xPc)
     local n_obs = size(xPc, 1)
     #is_dummy = isnan.(CA.getdata(xPc))
@@ -21,10 +64,6 @@ f_global = function(θsc, θgc, xPc)
     y = a1 .+  log.(a2) .* abs2.(cos.(b .- 0.2)) .* xPc.^2
 end
 
-() -> begin
-    include("test/test_scratch.jl")
-end
-
 @testset "PBMPopulationGlobalApplicator" begin
     n_obs = 3
     n_site = 5
@@ -47,5 +86,10 @@ end
     gr = Zygote.gradient((θP, θMs) -> sum(g(θP, θMs, xPc_NaN)), CA.getdata(θP), CA.getdata(θMs))
     @test all(isfinite.(gr[1])) # \thetaP
     @test all(isfinite.(gr[2])) # solves finite gradient for a2 for first site
+
+    n_site_m = 4
+    gm = create_nsite_applicator(g, n_site_m)
+    retm = gm(θP, θMs[1:n_site_m,:], xPc[:,1:n_site_m])
+    @test retm ≈ y_obs[:,1:n_site_m]
 end;