Made code faster

stecrotti · Apr 16, 2024 · 393ca39 · 393ca39
1 parent 45f14f1
commit 393ca39
Show file tree

Hide file tree

Showing 6 changed files with 12,437 additions and 12,435 deletions.
diff --git a/notebooks/sis_heterogeneous_compare_homogeneous.ipynb b/notebooks/sis_heterogeneous_compare_homogeneous.ipynb
diff --git a/src/MatrixProductBP.jl b/src/MatrixProductBP.jl
@@ -2,15 +2,15 @@ module MatrixProductBP
 
 using InvertedIndices: Not
 using ProgressMeter: Progress, ProgressUnknown, next!
-using TensorCast: @reduce, @cast, TensorCast 
+using TensorCast: @reduce, @cast, TensorCast
 using LoopVectorization
 using Tullio: @tullio
 using IndexedGraphs: IndexedGraphs, AbstractIndexedDiGraph, IndexedGraph, IndexedBiDiGraph,
     nv, ne, edges, vertices, inedges, outedges, src, dst, idx, neighbors, IndexedEdge,
     issymmetric
 using UnPack: @unpack
 using Random: shuffle!, AbstractRNG, GLOBAL_RNG
-using SparseArrays: rowvals, nonzeros, nzrange
+using SparseArrays: rowvals, nonzeros, nzrange, sparse
 using Distributions: Distributions, sample, sample!
 using Measurements: Measurement, ±
 using Statistics: mean, std

diff --git a/src/Models/Models.jl b/src/Models/Models.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, SparseMatrixCSC
+import SparseArrays: nonzeros, nzrange, rowvals, Symmetric, SparseMatrixCSC, sparse
 import TensorCast: @reduce, @cast, TensorCast 
 import ProgressMeter: Progress, next!, ProgressUnknown
 import LogExpFunctions: xlogx, xlogy

diff --git a/src/Models/epidemics/sis_heterogeneous.jl b/src/Models/epidemics/sis_heterogeneous.jl
@@ -1,6 +1,6 @@
 struct SIS_heterogeneous{T, N, F<:Real}
     g  :: IndexedGraph
-    λ  :: SparseMatrixCSC{F,Int64}
+    λ  :: Vector{F}
     ρ  :: Vector{F}
     ϕ  :: Vector{Vector{Vector{F}}}  # site observations
     ψ  :: Vector{Vector{Matrix{F}}}  # edge observations
@@ -17,7 +17,7 @@ struct SIS_heterogeneous{T, N, F<:Real}
         @assert all(length(ϕᵢ) == T + 1 for ϕᵢ in ϕ)
         @assert length(ψ) == 2*ne(g)
         @assert all(length(ψᵢⱼ) == T + 1 for ψᵢⱼ in ψ)
-        new{T,N,F}(g, λ, ρ, ϕ, ψ)
+        new{T,N,F}(g, nonzeros(λ), ρ, ϕ, ψ)
     end
 end
 
@@ -37,5 +37,5 @@ function SIS_heterogeneous(λ::SparseMatrixCSC{F,Int64}, ρ::Vector{F}, T::Int;
 end
 
 function sis_heterogeneous_factors(sis::SIS_heterogeneous{T,N,F}) where {T,N,F}
-    [fill(SIS_heterogeneousFactor([sis.λ[j,i] for j in sis.λ.rowval[nzrange(sis.λ,i)]], sis.ρ[i]), T + 1) for i in vertices(sis.g)]
+    [fill(SIS_heterogeneousFactor(sis.λ[nzrange(sis.g.A,i)], sis.ρ[i]), T + 1) for i in vertices(sis.g)]
 end
diff --git a/test/Project.toml b/test/Project.toml
@@ -4,5 +4,6 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 IndexedGraphs = "8a731c18-cfb7-4915-927e-cc26b56b67cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 TensorTrains = "89893e69-996d-40b1-ba32-8ff5f34c0dd5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,6 +1,6 @@
 using Test, MatrixProductBP
 using MatrixProductBP.Models
-using Graphs, IndexedGraphs, Random, Distributions
+using Graphs, IndexedGraphs, Random, Distributions, SparseArrays
 using TensorTrains
 using Aqua
 
@@ -19,6 +19,7 @@ include("pair_observations.jl")
 include("periodic.jl")
 include("sampling.jl")
 include("sirs_small_tree.jl")
+include("sis_heterogeneous_compare_homogeneous.jl")
 include("sis_infinite_graph.jl")
 include("sis_small_tree.jl")