Enhancement with sparse matrices

Author: FedericoFlorio

notebooks/sis_heterogeneous_compare_homogeneous.ipynb

@@ -0,0 +1,70 @@
+using SparseArrays
+using MatrixProductBP, MatrixProductBP.Models
+using UnPack, Graphs, IndexedGraphs, Plots, Statistics, ProgressMeter, Random
+ProgressMeter.ijulia_behavior(:clear)
+using JLD2
+import Measurements: value
+
+
+T = 3
+N = 5
+
+A = [0 1 1 0 0;
+     1 0 1 0 0;
+     1 1 0 1 0;
+     0 0 1 0 1;
+     0 0 0 1 0]
+g = IndexedGraph(A)
+
+λ_unif = 0.15
+ρ_unif = 0.12
+λ = λ_unif .* A
+λ = sparse(λ)
+ρ = fill(ρ_unif,N)
+γ = 0.13
+
+sis_u = SIS(g, λ_unif, ρ_unif, T; γ)
+sis_h = SIS_heterogeneous(λ, ρ, T; γ)
+
+
+svd_trunc = TruncBond(3)
+
+bp_u = mpbp(sis_u)
+iters_u, cb_u = iterate!(bp_u, maxiter=200; svd_trunc, tol=1e-12)
+b_bp_u = beliefs(bp_u)
+p_bp_u = [[bᵗ[INFECTIOUS] for bᵗ in bb] for bb in b_bp_u]
+
+bp_h = mpbp(sis_h)
+iters_h, cb_h = iterate!(bp_h, maxiter=200; svd_trunc, tol=1e-12)
+b_bp_h = beliefs(bp_h)
+p_bp_h = [[bᵗ[INFECTIOUS] for bᵗ in bb] for bb in b_bp_h]
+
+err = maximum([(p_bp_h[i][j]-p_bp_u[i][j])/p_bp_u[i][j] for i in eachindex(p_bp_u) for j in eachindex(p_bp_u[i])])
+
+println(err ≈ 0)
+
+
+A_ = copy(A)
+A_[3,2] = 0
+A_[4,3] = 0
+λ_ = copy(λ)
+λ_[3,2] = 0.0
+λ_[4,3] = 0.0
+
+sis_u_ = SIS(g, λ_unif, ρ_unif, T; γ)
+sis_h_ = SIS_heterogeneous(λ, ρ, T; γ)
+
+
+bp_u_ = mpbp(sis_u_)
+iters_u_, cb_u_ = iterate!(bp_u_, maxiter=200; svd_trunc, tol=1e-12)
+b_bp_u_ = beliefs(bp_u_)
+p_bp_u_ = [[bᵗ[INFECTIOUS] for bᵗ in bb] for bb in b_bp_u_]
+
+bp_h_ = mpbp(sis_h_)
+iters_h_, cb_h_ = iterate!(bp_h_, maxiter=200; svd_trunc, tol=1e-12)
+b_bp_h_ = beliefs(bp_h_)
+p_bp_h_ = [[bᵗ[INFECTIOUS] for bᵗ in bb] for bb in b_bp_h_]
+
+err_ = maximum([(p_bp_h_[i][j]-p_bp_u_[i][j])/p_bp_u_[i][j] for i in eachindex(p_bp_u_) for j in eachindex(p_bp_u_[i])])
+
+println(err_ ≈ 0)

notebooks/test.jl

@@ -11,7 +11,7 @@ using MatrixProductBP
 import IndexedGraphs: IndexedGraph, IndexedBiDiGraph, AbstractIndexedDiGraph, ne, nv, 
     outedges, idx, src, dst, inedges, neighbors, edges, vertices, IndexedEdge
 import UnPack: @unpack
-import SparseArrays: nonzeros, nzrange, rowvals, Symmetric
+import SparseArrays: nonzeros, nzrange, rowvals, Symmetric, SparseMatrixCSC
 import TensorCast: @reduce, @cast, TensorCast 
 import ProgressMeter: Progress, next!, ProgressUnknown
 import LogExpFunctions: xlogx, xlogy

src/Models/Models.jl

@@ -1,15 +1,15 @@
 struct SIS_heterogeneous{T, N, F<:Real}
     g  :: IndexedGraph
-    λ  :: Matrix{F}
+    λ  :: SparseMatrixCSC{F,Int64}
     ρ  :: Vector{F}
     ϕ  :: Vector{Vector{Vector{F}}}  # site observations
     ψ  :: Vector{Vector{Matrix{F}}}  # edge observations
-    function SIS_heterogeneous(g::IndexedGraph, λ::Matrix{F}, ρ::Vector{F},
+    function SIS_heterogeneous(g::IndexedGraph, λ::SparseMatrixCSC{F,Int64}, ρ::Vector{F},
             ϕ::Vector{Vector{Vector{F}}},
             ψ::Vector{Vector{Matrix{F}}}) where {F<:Real}
         @assert size(λ)[1] == size(λ)[2] == nv(g)
         @assert length(ρ) == nv(g)
-        @assert all(0 ≤ λᵢⱼ ≤ 1 for λᵢⱼ in λ)
+        @assert all(0 ≤ λᵢⱼ ≤ 1 for λᵢⱼ in !iszero(λ))
         @assert all(0 ≤ ρᵢ ≤ 1 for ρᵢ in ρ)
         N = nv(g)
         @assert length(ϕ) == N
@@ -21,13 +21,21 @@ struct SIS_heterogeneous{T, N, F<:Real}
     end
 end
 
-function SIS_heterogeneous(g::IndexedGraph{Int}, λ::Matrix{F}, ρ::Vector{F}, T::Int;
+function SIS_heterogeneous(g::IndexedGraph{Int}, λ::SparseMatrixCSC{F,Int64}, ρ::Vector{F}, T::Int;
         ψ = [[ones(2,2) for t in 0:T] for _ in 1:2*ne(g)],
         γ = 0.5,
         ϕ = [[t == 0 ? (length(γ)==1 ? [1-γ, γ] : [1-γ[i],γ[i]]) : ones(2) for t in 0:T] for i in vertices(g)]) where {F<:Real}
+        
     return SIS_heterogeneous(g, λ, ρ, ϕ, ψ)
 end
 
+function SIS_heterogeneous(λ::SparseMatrixCSC{F,Int64}, ρ::Vector{F}, T::Int; γ=0.5) where {F<:Real}
+    A = ones(Int,size(λ)[1],size(λ)[2]) - iszero.(λ)
+    g = IndexedGraph(A)
+    
+    return SIS_heterogeneous(g, λ, ρ, T; γ=γ)
+end
+
 function sis_heterogeneous_factors(sis::SIS_heterogeneous{T,N,F}) where {T,N,F}
-    [fill(SIS_heterogeneousFactor([sis.λ[x,i] for x in eachindex(sis.λ[:,i]) if x!=i], sis.ρ[i]), T + 1) for i in vertices(sis.g)]
+    [fill(SIS_heterogeneousFactor([sis.λ[j,i] for j in sis.λ.rowval[nzrange(sis.λ,i)]], sis.ρ[i]), T + 1) for i in vertices(sis.g)]
 end

src/Models/epidemics/sis_heterogeneous.jl

@@ -1,6 +1,3 @@
-# import MatrixProductBP: mpbp, nstates, prob_y, prob_xy, prob_yy
-
-
 const SUSCEPTIBLE = 1 
 const INFECTIOUS = 2
 

src/Models/epidemics/sis_heterogeneous_bp.jl

@@ -1,24 +1,25 @@
 @testset "SIS heterogeneous compare homogeneous" begin
     T = 3
     N = 5
-    
+
     A = [0 1 1 0 0;
-         1 0 1 0 0;
-         1 1 0 1 0;
-         0 0 1 0 1;
-         0 0 0 1 0]
+        1 0 1 0 0;
+        1 1 0 1 0;
+        0 0 1 0 1;
+        0 0 0 1 0]
     g = IndexedGraph(A)
 
     λ_unif = 0.15
     ρ_unif = 0.12
     λ = λ_unif .* A
+    λ = sparse(λ)
     ρ = fill(ρ_unif,N)
     γ = 0.13
 
     sis_u = SIS(g, λ_unif, ρ_unif, T; γ)
-    sis_h = SIS_heterogeneous(g, λ, ρ, T; γ)
+    sis_h = SIS_heterogeneous(λ, ρ, T; γ)
+
 
-    
     svd_trunc = TruncBond(3)
 
     bp_u = mpbp(sis_u)
@@ -34,4 +35,31 @@
     err = maximum([(p_bp_h[i][j]-p_bp_u[i][j])/p_bp_u[i][j] for i in eachindex(p_bp_u) for j in eachindex(p_bp_u[i])])
 
     @test err ≈ 0
+
+
+    A_ = copy(A)
+    A_[3,2] = 0
+    A_[4,3] = 0
+    λ_ = copy(λ)
+    λ_[3,2] = 0.0
+    λ_[4,3] = 0.0
+
+    sis_u_ = SIS(g, λ_unif, ρ_unif, T; γ)
+    sis_h_ = SIS_heterogeneous(λ, ρ, T; γ)
+
+
+    bp_u_ = mpbp(sis_u_)
+    iters_u_, cb_u_ = iterate!(bp_u_, maxiter=200; svd_trunc, tol=1e-12)
+    b_bp_u_ = beliefs(bp_u_)
+    p_bp_u_ = [[bᵗ[INFECTIOUS] for bᵗ in bb] for bb in b_bp_u_]
+
+    bp_h_ = mpbp(sis_h_)
+    iters_h_, cb_h_ = iterate!(bp_h_, maxiter=200; svd_trunc, tol=1e-12)
+    b_bp_h_ = beliefs(bp_h_)
+    p_bp_h_ = [[bᵗ[INFECTIOUS] for bᵗ in bb] for bb in b_bp_h_]
+
+    err_ = maximum([(p_bp_h_[i][j]-p_bp_u_[i][j])/p_bp_u_[i][j] for i in eachindex(p_bp_u_) for j in eachindex(p_bp_u_[i])])
+
+    @test err_ ≈ 0
+
 end