[WIP] RAT-SPNs

.gitignore

@@ -6,3 +6,4 @@ experiments
 *.mat
 *.csv
 *.pdf
+Manifest.toml

Project.toml

@@ -0,0 +1,15 @@
+name = "SumProductNetworks"
+uuid = "5f6e642e-680c-5a9c-a175-7c23ed4da89e"
+version = "0.1.4"
+
+[deps]
+AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
+
+[compat]
+julia = ">= 1.0"

src/SumProductNetworks.jl

@@ -20,17 +20,22 @@ using SparseArrays
 using Random
 using Printf
 
+# include custom distributions
+include("distributions.jl")
+
 # include general implementations
 include("nodes.jl")
 include("nodeFunctions.jl")
 include("networkFunctions.jl")
-
-# include approach specific implementations
-include("bmiTest.jl")
-include("structureUtilities.jl")
+include("regiongraphs.jl")
 
 # include utilities
 include("utilityFunctions.jl")
 include("io.jl")
 
+# include approach specific implementations
+include("bmiTest.jl")
+include("structureUtilities.jl")
+include("ratspn.jl")
+
 end # module

src/distributions.jl

@@ -0,0 +1,32 @@
+export NormalInverseGamma
+
+## Gaussian with Normal Inverse Gamma prior ##
+struct NormalInverseGamma{T<:Real} <: Distribution{Multivariate,Continuous}
+    μ::T
+    ν::T
+    a::T
+    b::T
+end
+NormalInverseGamma() = NormalInverseGamma(0.0, 1.0, 1.0, 1.0)
+
+Distributions.length(d::NormalInverseGamma) = 2
+
+function Distributions.logpdf(d::NormalInverseGamma, μ::T, σ²::T) where {T}
+    lp = _invgammalogpdf(d.a, d.b, σ²)
+    lp += normlogpdf(d.μ, sqrt(σ²) / d.ν, μ)
+    return l
+end
+
+@inline _invgammalogpdf(a::T, b::T, x::T) where {T<:Real} = a*log(b)-lgamma(a)-(a+one(T))*log(x)-b/x
+
+function Distributions._rand!(rng::AbstractRNG, d::NormalInverseGamma, out::AbstractArray{T,2}) where {T}
+    for p in eachslice(out, dims=[2])
+        @inbounds begin
+            p[2] = rand(InverseGamma(d.a, d.b))
+            p[1] = rand(Normal(d.μ, sqrt(p[2] / d.ν )))
+        end
+    end
+    return out
+end
+
+@inline Distributions.params(d::NormalInverseGamma) = (d.μ, d.ν, d.a, d.b)

src/ratspn.jl

@@ -0,0 +1,103 @@
+export ratspn
+#export templateLeaves, templatePartition, templateRegion
+
+function templateLeaves(likelihoods::AbstractVector{<:Distribution},
+                        priors::AbstractVector{<:Distribution},
+                        N::Int, D::Int, K::Int)
+    scopeVec = zeros(Bool, D)
+    parameters = map(prior -> rand(prior,K), priors)
+
+    return FactorizedDistributionGraphNode(
+                                    gensym("fact"),
+                                    scopeVec,
+                                    likelihoods,
+                                    priors,
+                                    parameters
+                                   )
+end
+
+function templatePartition(likelihoods::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_)
+
+    return PartitionGraphNode(
+                              gensym("partition"),
+                              scopeVec,
+                              obsVec,
+                              prior,
+                              children
+                             )
+end
+
+function buildRegion(llhs::AbstractVector, priors::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 -> buildLeaves(likelihoods, priors, N, D, K), 1:J)
+    else
+        map(k -> buildPartition(likelihoods, priors, N, D, K_sum, K_prod, J, K, depth+1, maxdepth), 1:J)
+    end
+
+    Ch = sum(length.(children))
+    scopeVec = zeros(Bool, D)
+    logweights = rand(Dirichlet(Ch, 1.0), K_)
+    active = zeros(Bool, size(logweights)...)
+    @assert size(logweights) == (Ch, K_)
+
+    return RegionGraphNode( gensym("region"), scopeVec, logweights, children )
+end
+
+
+function ratspn(x::AbstractMatrix{T};
+               Ksplits::Int=2,
+               Kparts::Int=2,
+               Ksums::Int=2,
+               Kdists::Int=5,
+               maxdepth::Int=2
+              ) where {T<:Real}
+
+    N,D = size(x)
+    isdiscrete = map(d -> all(isinteger, x[:,d]), 1:D)
+
+    K = map(d -> isdiscrete[d] ? length(unique(x[:,d])) : Inf, 1:D)
+
+    llhs = map(d -> isdiscrete[d] ? Categorical(K[d]) : Normal(), 1:D)
+    priors = map(d -> isdiscrete[d] ? Dirichlet(K[d], 1.0) : NormalInverseGamma(), 1:D)
+
+    return ratspn(N,D, llhs, priors, Ksplits, Kparts, Ksums, Kdists, maxdepth)
+end
+
+function ratspn(N::Int, D::Int,
+                llhs::AbstractVector{<:Distribution},
+                priors::AbstractVector{<:Distribution},
+                Ksplits::Int, # Number of partitions under each region
+                Kparts::Int, # Number of sub-regions under each partition
+                Ksums::Int, # Number of sum nodes per region
+                Kdists::Int, # Number of distibutions per terminal region
+                maxdepth::Int # Maximum number of consecutive region-partition pairs
+                )
+
+    # sanity checks
+    @assert length(llhs) == D == length(priors)
+
+
+
+    #return templateRegion()
+end