/
diffusion.jl
101 lines (86 loc) · 3.62 KB
/
diffusion.jl
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
"""
diffusion(g, p, n)
Run diffusion simulation on `g` for `n` steps with spread
probabilities based on `p`. Return a vector with the set of new
vertices reached at each step of the simulation.
### Optional Arguments
- `initial_infections=sample(vertices(g), 1)`: A list of vertices that
are infected at the start of the simulation.
- `watch=Vector()`: While simulation is always run on the full graph,
specifying `watch` limits reporting to a specific set of vertices reached
during the simulation. If left empty, all vertices will be watched.
- `normalize=false`: if `false`, set the probability of spread from a vertex ``i`` to
each of the outneighbors of ``i`` to ``p``. If `true`, set the probability of spread
from a vertex ``i`` to each of the `outneighbors` of ``i`` to
``\\frac{p}{outdegreee(g, i)}``.
- `rng=nothing`: A random generator to sample from.
"""
function diffusion(g::AbstractGraph{T},
p::Real,
n::Integer;
watch::AbstractVector=Vector{Int}(),
normalize::Bool=false,
rng::Union{Nothing, AbstractRNG} = nothing,
seed::Union{Nothing, Integer} = nothing,
initial_infections::AbstractVector=Graphs.sample(vertices(g), 1, rng=rng, seed=seed),
) where T
# Initialize
rng = rng_from_rng_or_seed(rng, seed)
watch_set = Set{T}(watch)
infected_vertices = BitSet(initial_infections)
vertices_per_step::Vector{Vector{T}} = [Vector{T}() for i in 1:n]
# Record initial infection
if !isempty(watch_set)
watched_initial_infections = intersect(watch_set, initial_infections)
vertices_per_step[1] = T.(collect(watched_initial_infections))
else
vertices_per_step[1] = T.(initial_infections)
end
# Run simulation
randsubseq_buf = zeros(T, Δout(g))
for step in 2:n
new_infections = Set{T}()
for i in infected_vertices
outn = outneighbors(g, i)
outd = length(outn)
if outd > 0
if normalize
local_p = p / outdegree(g, i)
else
local_p = p
end
randsubseq!(rng, randsubseq_buf, outn, local_p)
union!(new_infections, randsubseq_buf)
end
end
# Record only new infections
setdiff!(new_infections, infected_vertices)
if !isempty(watch_set)
vertices_per_step[step] = T.(collect(intersect(new_infections, watch_set)))
else
vertices_per_step[step] = collect(new_infections)
end
# Add new to master set of infected
union!(infected_vertices, new_infections)
end
return vertices_per_step
end
"""
diffusion_rate(results)
diffusion_rate(g, p, n; ...)
Given the results of a `diffusion` output or the parameters
to the `diffusion` simulation itself, (run and) return the rate of
diffusion as a vector representing the cumulative number of vertices
infected at each simulation step, restricted to vertices included
in `watch`, if specified.
"""
diffusion_rate(x::Vector{Vector{T}}) where T <: Integer = cumsum(length.(x))
diffusion_rate(g::AbstractGraph, p::Real, n::Integer;
watch::AbstractVector=Vector{Int}(),
normalize::Bool=false,
rng::Union{Nothing, AbstractRNG} = nothing,
seed::Union{Nothing, Integer} = nothing,
initial_infections::AbstractVector=Graphs.sample(vertices(g), 1, rng=rng, seed=seed),
) = diffusion_rate(diffusion(g, p, n,
initial_infections=initial_infections,
watch=watch, normalize=normalize, rng=rng, seed=seed))