1- function sample_ζresid_norm (app:: MeanHVIApproximation , z:: AbstractMatrix , ζP:: TP , ζMs:: TM ,
1+ function sample_ζresid_norm (app:: MeanHVIApproximation ,
2+ zP:: AbstractMatrix , zMs:: AbstractMatrix , ζP:: TP , ζMs:: TM ,
23 ϕq:: AbstractVector ;
34 int_ϕq= get_concrete (ComponentArrayInterpreter (ϕq)),
45 cor_ends
56) where {T,TP<: AbstractVector{T} ,TM<: AbstractMatrix{T} }
67 ϕuncc = int_ϕq (CA. getdata (ϕq))
7- n_θP, n_θMs, (n_θM, n_batch) = length (ζP), length (ζMs), size (ζMs)
8+ n_θP, n_θMs, (n_θM, n_batch), n_MC = length (ζP), length (ζMs), size (ζMs), size (zP, 1 )
89 # do not create a UpperTriangular Matrix of an AbstractGÜUArray in transformU_cholesky1
910 ρsP = isempty (ϕuncc[Val (:ρsP )]) ? similar (ϕuncc[Val (:ρsP )]) : ϕuncc[Val (:ρsP )] # required by zygote
1011 UPs, rangesP = transformU_blocks_cholesky1 (ρsP, cor_ends. P)
11- UP = transformU_block_cholesky1 (ρsP, cor_ends. P)
12- ρsM, rangesM = isempty (ϕuncc[Val (:ρsM )]) ? similar (ϕuncc[Val (:ρsM )]) : ϕuncc[Val (:ρsM )] # required by zygote
12+ ρsM = isempty (ϕuncc[Val (:ρsM )]) ? similar (ϕuncc[Val (:ρsM )]) : ϕuncc[Val (:ρsM )] # required by zygote
1313 # cholesky factor of the correlation: diag(UM' * UM) .== 1
1414 # coefficients ρsM can be larger than 1, still yielding correlations <1 in UM' * UM
15- UMs = transformU_block_cholesky1 (ρsM, cor_ends. M)
15+ UMs, rangesM = transformU_blocks_cholesky1 (ρsM, cor_ends. M)
1616 cf = ϕuncc[Val (:coef_logσ2_ζMs )]
1717 logσ2_logMs = vec (cf[1 , :] .+ cf[2 , :] .* ζMs)
1818 logσ2_ζP = vec (CA. getdata (ϕuncc[Val (:logσ2_ζP )]))
1919 # CUDA cannot multiply BlockDiagonal * Diagonal, construct already those blocks
2020 σMs = reshape (exp .(logσ2_logMs ./ 2 ), n_θM, :)
2121 σP = exp .(logσ2_ζP ./ 2 )
22- # create random numbers from U diag(σ) z = (σ .* z)
23-
24-
25- # # BlockDiagonal does work with CUDA, but not with combination of Zygote and CUDA
26- # # need to construct full matrix for CUDA
27- # Uσ, diagUσ = _compute_choleskyfactor(UP, UM, σP, σMs, n_batch) # inferred only BlockDiagonal
28- # #diagUσ = diag(Uσ)::typeof(σP) # elements of the diagonal: standard deviations
29- # n_MC = size(urandn, 2) # TODO transform urandn
30- # # is this multiplication efficient if Uσ is not concrete but only sumtype BlockDiagonal?
31- # ζ_resids_parfirst = (Uσ' * urandn) #::typeof(urandn) # n_par x n_MC
32- # #ζ_resids_parfirst = urandn' * Uσ # n_MC x n_par
33- # # need to handle empty(ζP) explicitly, otherwise Zygote tries to take gradient
34- # ζP_resids = isempty(ζP) ? ζ_resids_parfirst[1:0, :] : ζ_resids_parfirst[1:n_θP, :]
35- # ζMs_parfirst_resids = reshape(ζ_resids_parfirst[(n_θP+1):end, :], n_θM, n_batch, n_MC)
36-
22+ # create random numbers from U diag(σ) z = U (σ .* z)
23+ # for each block separately
24+ # σzMs = σMs .* reshape(zMs, (n_θM, n_batch, n_MC))
25+ # Ui, ri = first(zip(UMs, rangesM))
26+ zMs_subjects = reshape (zMs, (n_MC, n_θM, n_batch))
27+ # σM, zM = first(zip(eachcol(σMs), eachslice(zMs_subjects; dims=3)))
28+ cat3 = (x,y) -> cat (x,y,dims= 3 )
29+ # map across subjects (n_batch)
30+ # ζMs_vec = map(eachcol(σMs), eachslice(zMs_subjects; dims=3)) do σM, zM
31+ fBlock = let UMs = UMs, rangesM = collect (rangesM) # without collect, type unstable
32+ function fBlock_inner (σM, zM)
33+ mapreduce (vcat, UMs, rangesM) do Ui, ri # n_θM, n_MC
34+ ri:: UnitRange{Int}
35+ # (zM[:,ri]' * Ui * diagm(σM[ri]))'
36+ diagm (σM[ri]) * Ui' * zM[:, ri]'
37+ end
38+ end
39+ end
40+ ζMs1 = ChainRulesCore. @ignore_derivatives fBlock (σMs[:,1 ], zMs_subjects[:,:,1 ])
41+ # TB = Base.infer_return_type(fBlock, (typeof(σMs[:,1]), typeof(zMs_subjects[:,:,1])))
42+ # ζMs_vec = map(fBlock, eachcol(σMs[:,2:end]), eachslice(zMs_subjects[:,:,2:end]; dims=3), init = ζMs1)
43+ ζMs_vec = map ((σM, zM) -> fBlock (σM, zM):: typeof (ζMs1), eachcol (σMs), eachslice (zMs_subjects; dims= 3 ))
44+ # zM = zMs_subjects[:,:,1]
45+ # ζMs_vec = [fBlock(σMs[:,i], zM) for i in axes(σMs, 2)]
46+ # ζMs_vec = [fBlock(σMs[:,1], zM)]
47+ # ζMs_vec = [fBlock(σMs[:,1], zM) for i in 1:n_batch]
48+ # ζMs_vec = [fBlock(σMs[:,1], zM) for i in axes(σMs, 2)]
49+ # ζMs_vec = [fBlock(σM, zM) for (σM, zM) in zip(eachcol(σMs), eachslice(zMs_subjects; dims=3))]
50+ # concatenate so that n_MC is last dimension
51+ local ζMs_parfirst_resids = stack (ζMs_vec; dims = 2 ) # n_θM, n_batch, n_MC
52+ # size(ζMs_parfirst_resids)
3753
54+ # std(ζMs_parfirst_resids[1,1,:])
55+ # std(ζMs_parfirst_resids[1,end,:])
56+ # σzMs_stacked = reshape(σzMs, (n_θM, n_batch * n_MC))
57+ # ζMs_resids_stacked = mapreduce(vcat, UMs, rangesM) do Ui, ri
58+ # #Ui * σzMs_stacked[ri, :]
59+ # Uσ = Ui * σMs[ri,:]
60+ # Uσ' * zMs[:,ri]'
61+ # diagm(σMs[ri,:]) * Ui' * zMs[:,ri]'
62+ # end
63+ # ζMs_parfirst_resids = reshape(ζMs_resids_stacked, n_θM, n_batch, n_MC)
64+ #
65+ # σzP = σP .* reshape(zP, (n_θP, n_MC))
66+ # Ui, ri = first(zip(UPs, rangesP))
67+ local ζP_resids = if isempty (ζP)
68+ # provide init of correct empty matrix type
69+ # single branch with init to mapreduce has Problems with Zygote
70+ ChainRulesCore. @ignore_derivatives ζMs_parfirst_resids[1 : 0 ,1 ,1 : n_MC]
71+ else
72+ mapreduce (vcat, UPs, rangesP) do Ui, ri
73+ diagm (σP[ri]) * Ui' * zP[:,ri]'
74+ end
75+ end
76+ #
77+ () -> begin
78+ tmp = tmp1 = (σzP[ri,:]' * Ui)'
79+ tmp = tmp2 = (zP * (Ui * diagm (σP)))'
80+ tmp = tmp3 = (zP * Ui * diagm (σP))' # σ on wrong side of U
81+ tmp = tmp4 = diagm (σP) * Ui' * zP' # σ on wrong side of U
82+ std (tmp[2 ,:])
83+ end
84+ diagUσ = vcat (σP, vec (σMs))
3885 ζP_resids, ζMs_parfirst_resids, diagUσ
39- # #map(std, eachcol(ζ_resids_parfirst[:, 3:8]))
40- # ζ_resid = transpose_mPMs_sitefirst(ζ_resids_parfirst; intm_PMs_parfirst)
41- # #map(std, eachcol(ζ_resid[:, 3:8])) # all ~ 0.1 in sample_ζresid_norm cpu
42- # #map(std, eachcol(ζ_resid[:, 2 + n_batch .+ (-1:5)])) # all ~ 100, except first two
43- # # returns AbstractGPUuArrays to either continue on GPU or need to transfer to CPU
44- # ζ_resid, diagUσ
4586end
0 commit comments