def add_timeseries(self,
t,
S,
I,
R=None,
colordict=None,
label=None,
**kwargs):
r'''This allows us to include some additional timeseries for comparision
with the simulation. So for example, if we perform a simulation and
want to plot the simulation but also a prediction, this is what we
would use.
:Arguments:
**t** list
the times
**S** list
the number susceptible at each time
**I** list
the number infected at each time
**R** list (default None)
the number recovered at each time
**colordict** dict (default None)
a dictionary mapping 'S', 'I', and (if SIR) 'R' to the color
desired for their plots. Defaults to the same as the simulation
**label** (string)
The label to be used for these plots in the legend.
****kwargs**
any matplotlib key word args to affect how the curve is shown.
:Returns:
**ts** timeseries object
:Modifies:
This adds the timeseries object `ts` to the internal _time_series_list_
'''
if (R is not None and not self.SIR):
raise EoN.EoNError("cannot define R if SIR isn't True")
if (R is None and self.SIR):
raise EoN.EoNError("cannot have SIR True if no R defined")
if colordict is None:
colordict = self.colordict
ts = self._time_series_(t,
S,
I,
R,
self.SIR,
colordict=colordict,
label=label,
**kwargs)
self._time_series_list_.append(ts)
return ts
def sim_update_colordict(self, colordict):
r'''updates the colordict for the simulation
:Arguments:
**colordict** dict
the new colordict
'''
if not (colordict.has_key('S') and colordict.has_key('I')):
raise EoN.EoNError("colordict must have keys 'S' and 'I'")
if self.SIR and not colordict.has_key('R'):
raise EoN.EoNError("if SIR, then colordict must have key 'R'")
self.colordict = colordict
def update_ts_colordict(self, ts, colordict):
r'''updates the colordict for time series plots
:Arguments:
**ts** timeseries object
the timeseries object whose key word args we are updating.
**colordict** dict
the new colordict
'''
if not (colordict.has_key('S') and colordict.has_key('I')):
raise EoN.EoNError("colordict must have keys 'S' and 'I'")
if self.SIR and not colordict.has_key('R'):
raise EoN.EoNError("if SIR, then colordict must have key 'R'")
ts.colordict = colordict
def add_timeseries(self,
t,
S,
I,
R=None,
colordict=None,
label=None,
**kwargs):
r'''This allows us to include some additional timeseries for comparision
with the simulation. The most likely source of these would be an analytic
calculation.
Arguments :
t (list)
the times
S (list)
the number susceptible at each time
I (list)
the number infected at each time
R (list default None)
the number recovered at each time
colordict (dict)
a dictionary mapping 'S', 'I', and (if SIR) 'R' to the color
desired for their plots. Defaults to the same as the simulation
label (string)
The label to be used for these plots in the legend.
**kwargs
any matplotlib key word args to affect how the curve is shown.
Returns :
ts timeseries object
'''
if (R is not None and not self.SIR):
raise EoN.EoNError("cannot define R if SIR isn't True")
if (R is None and self.SIR):
raise EoN.EoNError("cannot have SIR True if no R defined")
if colordict is None:
colordict = self.colordict
ts = self._time_series_(t,
S,
I,
R,
self.SIR,
colordict=colordict,
label=label,
**kwargs)
self._time_series_list_.append(ts)
return ts
def transmission_tree(self):
r'''
Produces a MultiDigraph whose edges correspond to transmission events.
If SIR, then this is a tree (or a forest).
:Returns:
**T** a directed Multi graph
T has all the information in ``transmissions``.
An edge from u to v with time t means u transmitted to v at time t.
:Warning:
Although we refer to this as a "tree", if the disease is SIS, there
are likely to be cycles and/or repeated edges. If the disease is SIR
but there are multiple initial infections, then this will be a "forest".
If it's an SIR, then this is a tree (or forest).
The graph contains only those nodes that are infected at some point.
'''
if self._transmissions_ is None:
raise EoN.EoNError("transmissions were not provided when created")
T = nx.MultiDiGraph()
for t, u, v in self._transmissions_:
if u is not None:
T.add_edge(u, v, time=t)
return T
def sim_update_colordict(self, colordict):
r'''
CHANGES NEEDED. I think just make the test check that keys of colordict
and possible statuses match.
updates the colordict for the simulation
:Arguments:
**colordict** dict
the new colordict
'''
if not (colordict.has_key('S') and colordict.has_key('I')):
raise EoN.EoNError("colordict must have keys 'S' and 'I'")
if self.SIR and not colordict.has_key('R'):
raise EoN.EoNError("if SIR, then colordict must have key 'R'")
self.colordict = colordict
def __init__(self, t, S, I, R=None, colordict=None, label=None, **kwargs):
if colordict is None:
raise EoN.EoNError("colordict must be defined")
self._S_ = S
self._I_ = I
self._R_ = R
self._t_ = t
self.colordict = colordict
self.label=label
self.plt_kwargs = kwargs
def update_ts_colordict(self, ts, colordict):
r'''
CHANGES NEEDED. I think just make the test check that keys of colordict
and possible statuses match.
updates the colordict for time series plots
:Arguments:
**ts** timeseries object
the timeseries object whose key word args we are updating.
**colordict** dict
the new colordict
'''
if not (colordict.has_key('S') and colordict.has_key('I')):
raise EoN.EoNError("colordict must have keys 'S' and 'I'")
if self.SIR and not colordict.has_key('R'):
raise EoN.EoNError("if SIR, then colordict must have key 'R'")
ts.colordict=colordict
def R(self):
r'''
See notes for S
Returns the number recovered at each time
Generally better to get these all through summary()'''
if 'R' in self._possible_statuses_:
return self._summary_[1]['R']
else:
raise EoN.EoNError("'R' is not a possible status")
def legend(self, timeseries_plot_index = 0, loc='best'):
#timeseries plot can be:
#integer from 0 to len(timeseries)-1
# or string 'All'
if not self.timeseries_axi:
raise EoN.EoNError("no time series axes defined")
elif timeseries_plot_index == 'All':
for ax in self.timeseries_axi:
ax.legend(loc=loc)
else:
self.timeseries_axi[timeseries_plot_index].legend(loc=loc)
def transmissions(self):
r'''Returns a list of tuples (t,u,v) stating that node u infected node
v at time t. In the standard code, if v was already infected at tmin, then
the source is None
Note - this only includes successful transmissions. So if u tries
to infect v, but fails because v is already infected this is not
recorded.'''
if self._transmissions_ is None:
raise EoN.EoNError("transmissions were not provided when created")
return self._transmissions_
def S(self):
r'''
If ``'S'`` is a state, then this will return the number susceptible at each time.
Else it raises an error
Generally better to get these all through ``summary()`` '''
if 'S' in self._possible_statuses_:
return self._summary_[1]['S']
else:
raise EoN.EoNError("'S' is not a possible status")
def generate_network(Pk, N, ntries=100):
r'''Generates an N-node random network whose degree distribution is given by Pk'''
counter = 0
while counter < ntries:
counter += 1
ks = []
for ctr in range(N):
ks.append(Pk())
if sum(ks) % 2 == 0:
break
if sum(ks) % 2 == 1:
raise EoN.EoNError("cannot generate even degree sum")
G = nx.configuration_model(ks)
return G
def sim_update_color_dict(self, color_dict):
r'''
updates the color_dict for the simulation
:Arguments:
**color_dict** dict
the new color_dict
'''
for status in self._possible_statuses_:
if status not in color_dict:
raise EoN.EoNError("Status {} is not in color_dict".format(status))
self.sim_color_dict = color_dict
self._simulation_time_series_.color_dict=color_dict
def update_ts_color_dict(self, ts, color_dict):
r'''
updates the color_dict for time series plots
:Arguments:
**ts** timeseries object
the timeseries object whose key word args we are updating.
**color_dict** dict
the new color_dict
'''
for status in ts._D_:
if status not in color_dict:
raise EoN.EoNError("Status {} is not in color_dict".format(status))
ts.color_dict=color_dict
def __init__(self, G, node_history, transmissions, SIR = True, pos = None,
colordict=None, possible_statuses = None):
r'''
CHANGES PLANNED - remove SIR flag.
Make possible_statuses an optional entry where SIR used to be. If
not given, just make it a list of the values found in node_history.
Make clear that the new things replacing _S_ etc aren't defined until
summary is called.
Make transmissions optional
:Arguments:
**G** The graph
**node_history** (dict)
`node_history[node]` is a tuple (times, statuses) where
- `times` is a list of the times at which `node` changes status.
the first entry is the initial time.
- `statuses` is a list giving the status the new status of the node
at each corresponding time.
**transmissions** (list)
Each event which is induced by a neighbor appears (in order) in
`transmissions`. It appears as a triple `(time, source, target)`
where
- `time` is the time of the event
- `source` is the neighbor inducing the transition.
- `target` is the node undergoing the transition.
**SIR** Boolean. I plan to remove this.
**pos** (dict - default None)
The `pos` argument to be given to the networkx plotting commands.
**colordict** (dict - default None)
A dictionary stating for each status what color is to be used when
plotting.
If no colors are given, it will attempt to use a greenish color for
`'S'` and a reddish color for `'I'` and gray for `'R'`. These
should be color-blind friendly, despite appearing green/red to
me.
**possible_statuses** (ordered tuple or list)
The possible statues we expect to see.
'''
if SIR and possible_statuses is not None:
raise EoN.EoNError('if SIR, then possible_statuses should be None, (may change in future versions)')
if colordict is None:
if possible_statuses is None: #SIS or SIR
colordict = {'S':'#009a80','I':'#ff2000', 'R':'gray'}
else:
print("If you're seeing this message then I forgot to fix this."+
"Submit a bug report and say nasty things about my laziness." +
"if you want to plot you'll need to update the colordict" +
"see sim_update_colordict and update_ts_colordict")
colordict = {}#status: X for status in possible_statuses}
self.G = G
self._node_history_ = node_history
self._transmissions_ = transmissions
self.SIR = SIR
self.sim_colordict = colordict
self.pos = pos
self.summary() #defines self._t_, self._S_, self._I_, and self._R_
self._time_series_list_ = []
self._simulation_time_series_ = self._time_series_(self._t_, self._S_, self._I_, self._R_,
colordict=self.sim_colordict, label = 'Simulation')
self._time_series_list_.append(self._simulation_time_series_)
def subsample(report_times, times, status1, status2=None,
status3 = None):
r'''
Given
S, I, and/or R as lists of numbers of nodes of the given status
at given times
returns them
subsampled at specific report_times.
:Arguments:
**report_times** iterable (ordered)
times at which we want to know state of system
**times** iterable (ordered)
times at which we have the system state (assumed no change
between these times)
**status1** iterable
generally S, I, or R
number of nodes in given status at corresponding time in times.
**status2** iterable (optional, default None)
generally S, I, or R
number of nodes in given status at corresponding time in times.
**status3** iterable (optional, default None)
generally S, I, or R
number of nodes in given status at corresponding time in times.
:Returns:
If only status1 is defined
**report_status1** numpy array
gives ``status1`` subsampled just at ``report_times``.
If more are defined then it returns a list, either
**[report_status1, report_status2]**
or
**[report_status1, report_status2, report_status3]**
In each case, these are subsampled just at report_times.
:SAMPLE USE:
::
import networkx as nx
import EoN
import numpy as np
import matplotlib.pyplot as plt
""" in this example we will run 100 stochastic simulations.
Each simulation will produce output at a different set
of times. In order to calculate an average we will use
subsample to find the epidemic sizes at a specific set
of times given by report_times.
"""
G = nx.fast_gnp_random_graph(10000,0.001)
tau = 1.
gamma = 1.
report_times = np.linspace(0,5,101)
Ssum = np.zeros(len(report_times))
Isum = np.zeros(len(report_times))
Rsum = np.zeros(len(report_times))
iterations = 100
for counter in range(iterations):
t, S, I, R = EoN.fast_SIR(G, tau, gamma, initial_infecteds = range(10))
#t, S, I, and R have an entry for every single event.
newS, newI, newR = EoN.subsample(report_times, t, S, I, R)
#could also do: newI = EoN.subsample(report_times, t, I)
plt.plot(report_times, newS, linewidth=1, alpha = 0.4)
plt.plot(report_times, newI, linewidth=1, alpha = 0.4)
plt.plot(report_times, newR, linewidth=1, alpha = 0.4)
Ssum += newS
Isum += newI
Rsum += newR
Save = Ssum / float(iterations)
Iave = Isum / float(iterations)
Rave = Rsum / float(iterations)
plt.plot(report_times, Save, "--", linewidth = 5, label = "average")
plt.plot(report_times, Iave, "--", linewidth = 5)
plt.plot(report_times, Rave, "--", linewidth = 5)
plt.legend(loc = "upper right")
plt.savefig("tmp.pdf")
If only one of the sample times is given then returns just that.
If report_times goes longer than times, then this simply assumes the
system freezes in the final state.
This uses a recursive approach if multiple arguments are defined.
'''
if report_times[0] < times[0]:
raise EoN.EoNError("report_times[0]<times[0]")
report_status1 = []
next_report_index = 0
next_observation_index = 0
while next_report_index < len(report_times):
while next_observation_index < len(times) and \
times[next_observation_index]<= report_times[next_report_index]:
candidate = status1[next_observation_index]
next_observation_index += 1
report_status1.append(candidate)
next_report_index +=1
report_status1= np.array(report_status1)
if status2 is not None:
if status3 is not None:
report_status2, report_status3 = subsample(report_times, times, status2, status3)
return report_status1, report_status2, report_status3
else:
report_status2 = subsample(report_times, times, status2)
return report_status1, report_status2
else:
return report_status1
