BayesianSumProductNetworks/src/templateStructure.jl at master · trappmartin/BayesianSumProductNetworks · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
export templateLeaves, templatePartition, templateRegion

function templateLeaves(alpha_leaf_prior::Float64,
                        priors::AbstractVector{<:Distribution},
                        likelihoods::AbstractVector{<:Distribution},
                        sstats::AbstractVector{<:AbstractSufficientStats},
                        N::Int, D::Int, K::Int)
    ids = map(k -> gensym("factorization"), 1:K)
    scopeVec = zeros(Bool, D)
    obsVec = zeros(Bool, N, K)
    priors_ = deepcopy(priors)
    likelihoods_ = deepcopy(likelihoods)
    sstats_ = mapreduce(k -> deepcopy(sstats), hcat, 1:K)

    return FactorizedDistributionGraphNode(
                                    gensym("fact"),
                                    scopeVec,
                                    obsVec,
                                    priors_,
                                    likelihoods_,
                                    sstats_
                                   )
end

function templateLeaves(alpha_leaf_prior::Float64,
                        priors::Vector{<:AbstractVector{<:Distribution}},
                        likelihoods::Vector{<:AbstractVector{<:Distribution}},
                        sstats::Vector{<:AbstractVector{<:AbstractSufficientStats}},
                        N::Int, D::Int, K::Int)
    ids = map(k -> gensym("factorization"), 1:K)
    scopeVec = zeros(Bool, D)
    obsVec = zeros(Int, N, K)
    weightpriors = map(d -> Dirichlet(length(likelihoods[d]), alpha_leaf_prior), 1:D)
    logweights = map(d -> log.(rand(weightpriors[d], K)), 1:D)
    priors_ = deepcopy(priors)
    likelihoods_ = deepcopy(likelihoods)
    sstats_ = map(d -> mapreduce(k -> deepcopy(sstats[d]), hcat, 1:K), 1:D)

    return FactorizedMixtureGraphNode(
                                    gensym("factmixture"),
                                    scopeVec,
                                    obsVec,
                                    weightpriors,
                                    logweights,
                                    priors_,
                                    likelihoods_,
                                    sstats_
                                   )
end

function templatePartition(alpha_region_prior::Float64,
                           alpha_partition_prior::Float64,
                           alpha_leaf_prior::Float64,
                           priors_leaf::AbstractVector,
                           likelihoods::AbstractVector,
                           sstats::AbstractVector,
                           N::Int, D::Int,
                           K_sum::Int, K_prod::Int,
                           J::Int, K::Int,
                           depth::Int, maxdepth::Int)

    children = if depth == maxdepth
        map(k -> templateLeaves(alpha_leaf_prior, priors_leaf, likelihoods, sstats, N, D, K), 1:K_prod)
    else
        map(k -> templateRegion(alpha_region_prior, alpha_partition_prior, alpha_leaf_prior,
                                priors_leaf, likelihoods, sstats, N, D, K_sum, K_prod, J, K, depth+1, maxdepth), 1:K_prod)
    end

    K_ = mapreduce(child -> length(child), *, children)
    scopeVec = zeros(Bool, D)
    obsVec = zeros(Bool, N, K_)
    prior = Dirichlet(K_prod, alpha_partition_prior)

    return PartitionGraphNode(
                              gensym("partition"),
                              scopeVec,
                              obsVec,
                              prior,
                              children
                             )
end

function templateRegion(alpha_region_prior::Float64,
                        alpha_partition_prior::Float64,
                        alpha_leaf_prior::Float64,
                        priors_leaf::AbstractVector,
                        likelihoods::AbstractVector,
                        sstats::AbstractVector,
                        N::Int, D::Int,
                        K_sum::Int, K_prod::Int,
                        J::Int, K::Int,
                        depth::Int, maxdepth::Int; root = false)

    K_ = root ? 1 : K_sum
    children = if depth == maxdepth
        map(k -> templateLeaves(alpha_leaf_prior, priors_leaf, likelihoods, sstats, N, D, K), 1:J)
    else
        map(k -> templatePartition(alpha_region_prior, alpha_partition_prior, alpha_leaf_prior,
                                   priors_leaf, likelihoods, sstats, N, D, K_sum, K_prod, J, K, depth+1, maxdepth), 1:J)
    end

    Ch = sum(length.(children))
    scopeVec = zeros(Bool, D)
    obsVec = zeros(Bool, N, K_)
    prior = Dirichlet(Ch, alpha_region_prior)
    logweights = convert(Matrix, reshape(mapreduce(_ -> rand(prior), hcat, 1:K_), Ch, K_))
    active = zeros(Bool, size(logweights)...)
    @assert size(logweights) == (Ch, K_)

    return RegionGraphNode(
                           gensym("region"),
                            scopeVec,
                            obsVec,
                            logweights,
                            active,
                            prior,
                            children
                          )

end