def __init__(self, width, height):
super().__init__(width, height, "Shapes")
self.time = 0
self.batch = pyglet.graphics.Batch()
self.circle = shapes.Circle(360, 240, 100, color=(255, 225, 255), batch=self.batch)
self.circle.opacity = 127
# Rectangle with center as anchor
self.square = shapes.Rectangle(360, 240, 200, 200, color=(55, 55, 255), batch=self.batch)
self.square.anchor_x = 100
self.square.anchor_y = 100
# Large transparent rectangle
self.rectangle = shapes.Rectangle(0, 190, 720, 100, color=(255, 22, 20), batch=self.batch)
self.rectangle.opacity = 64
self.line = shapes.Line(0, 0, 0, 480, width=4, color=(200, 20, 20), batch=self.batch)
self.triangle = shapes.Triangle(10, 10, 190, 10, 100, 150, color=(10, 255, 10), batch=self.batch)
self.triangle.opacity = 175
self.arc = shapes.Arc(50, 300, radius=40, segments=25, angle=4, color=(255, 255, 255), batch=self.batch)
def visualize(
model,
network,
sampling_dt,
ignore_plot_compartments=[],
quarantine_compartments=[],
config=None,
):
"""
Start a visualization of a stochastic simulation.
Parameters
==========
model : epipack.stochastic_epi_models.StochasticEpiModel
An initialized StochasticEpiModel.
network: dict
A stylized network in the netwulf-format
(see https://netwulf.readthedocs.io/en/latest/python_api/post_back.html)
where instead of 'x' and 'y', node positions are saved in 'x_canvas'
and 'y_canvas'. Example:
.. code:: python
stylized_network = {
"xlim": [0, 833],
"ylim": [0, 833],
"linkAlpha": 0.5,
"nodeStrokeWidth": 0.75,
"links": [
{"source": 0, "target": 1, "width": 3.0 }
],
"nodes": [
{"id": 0,
"x_canvas": 436.0933431058901,
"y_canvas": 431.72418500564186,
"radius": 20
},
{"id": 1,
"x_canvas": 404.62184898400426,
"y_canvas": 394.8158724310507,
"radius": 20
}
]
}
sampling_dt : float
The amount of simulation time that's supposed to pass
with a single update.
ignore_plot_compartments : list, default = []
List of compartment objects that are supposed to be
ignored when plotted.
quarantine_compartments : list, default = []
List of compartment objects that are supposed to be
resemble quarantine (i.e. temporarily
losing all connections)
config : dict, default = None
A dictionary containing all possible configuration
options. Entries in this dictionary will overwrite
the default config which is
.. code:: python
_default_config = {
'plot_sampled_curve': True,
'draw_links':True,
'draw_nodes':True,
'n_circle_segments':16,
'plot_height':120,
'bgcolor':'#253237',
'curve_stroke_width':4.0,
'node_stroke_width':1.0,
'link_color': '#4b5a62',
'node_stroke_color':'#000000',
'node_color':'#264653',
'bound_increase_factor':1.0,
'update_dt':0.04,
'show_curves':True,
'draw_nodes_as_rectangles':False,
'show_legend': True,
'legend_font_color':'#fafaef',
'legend_font_size':10,
'padding':10,
'compartment_colors':_colors
}
"""
# update the config and compute some helper variables
cfg = deepcopy(_default_config)
if config is not None:
cfg.update(config)
palette = cfg['palette']
if type(palette) == str:
if 'link_color' not in cfg:
cfg['link_color'] = col.hex_link_colors[palette]
if 'bgcolor' not in cfg:
cfg['bgcolor'] = col.hex_bg_colors[palette]
if 'compartment_colors' not in cfg:
cfg['compartment_colors'] = [
col.colors[this_color] for this_color in col.palettes[palette]
]
bgcolor = [_ / 255
for _ in list(bytes.fromhex(cfg['bgcolor'][1:]))] + [1.0]
bgY = 0.2126 * bgcolor[0] + 0.7152 * bgcolor[1] + 0.0722 * bgcolor[2]
if cfg['legend_font_color'] is None:
if bgY < 0.5:
cfg['legend_font_color'] = '#fafaef'
else:
cfg['legend_font_color'] = '#232323'
width = network['xlim'][1] - network['xlim'][0]
height = network['ylim'][1] - network['ylim'][0]
with_plot = cfg['show_curves'] and set(ignore_plot_compartments) != set(
model.compartments)
if with_plot:
height += cfg['plot_height']
plot_width = width
plot_height = cfg['plot_height']
else:
plot_height = 0
with_legend = cfg['show_legend']
if with_legend:
legend_batch = pyglet.graphics.Batch()
#x, y = legend.get_location()
#legend.set_location(x - width, y)
# create a test label to get the actual dimensions
test_label = pyglet.text.Label('Ag')
dy = test_label.content_height * 1.1
del (test_label)
legend_circle_radius = dy / 2 / 2
distance_between_circle_and_label = 2 * legend_circle_radius
legend_height = len(model.compartments) * dy + cfg['padding']
# if legend is shown in concurrence to the plot,
# move the legend to be on the right hand side of the plot,
# accordingly make the plot at least as tall as
# the demanded height or the legend height
if with_plot:
plot_height = max(plot_height, legend_height)
legend_y_offset = legend_height
max_text_width = 0
legend_objects = [
] # this is a hack so that the garbage collector doesn't delete our stuff
for iC, C in enumerate(model.compartments):
this_y = legend_y_offset - iC * dy - cfg['padding']
this_x = width + cfg['padding'] + legend_circle_radius
label = pyglet.text.Label(
str(C),
font_name=('Helvetica', 'Arial', 'Sans'),
font_size=cfg['legend_font_size'],
x=this_x + legend_circle_radius +
distance_between_circle_and_label,
y=this_y,
anchor_x='left',
anchor_y='top',
color=list(bytes.fromhex(cfg['legend_font_color'][1:])) +
[255],
batch=legend_batch)
legend_objects.append(label)
#if not cfg['draw_nodes_as_rectangles']:
if True:
disk = shapes.Circle(
this_x,
this_y - (dy - 1.25 * legend_circle_radius) / 2,
legend_circle_radius,
segments=64,
color=cfg['compartment_colors'][iC],
batch=legend_batch,
)
circle = shapes.Arc(
this_x,
this_y - (dy - 1.25 * legend_circle_radius) / 2,
legend_circle_radius,
segments=64 + 1,
color=list(bytes.fromhex(cfg['legend_font_color'][1:])),
batch=legend_batch,
)
legend_objects.extend([disk, circle])
#else:
# rect = shapes.Rectangle(this_x,
# this_y - (dy-1.5*legend_circle_radius)/2,
# 2*legend_circle_radius,
# 2*legend_circle_radius,
# color = _colors[iC],
# batch=legend_batch,
# )
# legend_objects.append(rect)
max_text_width = max(max_text_width, label.content_width)
legend_width = 2*cfg['padding'] \
+ 2*legend_circle_radius \
+ distance_between_circle_and_label \
+ max_text_width
# if legend is shown in concurrence to the plot,
# move the legend to be on the right hand side of the plot,
# accordingly make the plot narrower and place the legend
# directly under the square network plot.
# if not, make the window wider and show the legend on
# the right hand side of the network plot.
if with_plot:
for obj in legend_objects:
obj.x -= legend_width
plot_width = width - legend_width
else:
width += legend_width
size = (width, height)
# overwrite network style with the epipack default style
network['linkColor'] = cfg['link_color']
network['nodeStrokeColor'] = cfg['node_stroke_color']
for node in network['nodes']:
node['color'] = cfg['node_color']
N = len(network['nodes'])
# get the OpenGL shape objects that comprise the network
network_batch = get_network_batch(
network,
yoffset=plot_height,
draw_links=cfg['draw_links'],
draw_nodes=cfg['draw_nodes'],
draw_nodes_as_rectangles=cfg['draw_nodes_as_rectangles'],
n_circle_segments=cfg['n_circle_segments'],
)
lines = network_batch['lines']
disks = network_batch['disks']
circles = network_batch['circles']
node_to_lines = network_batch['node_to_lines']
batch = network_batch['batch']
# initialize a simulation state that has to passed to the app
# so the app can change simulation parameters
simstate = SimulationStatus(len(network['nodes']), sampling_dt)
# intialize app
window = App(*size, simulation_status=simstate, resizable=True)
glClearColor(*bgcolor)
# handle different strokewidths
if 'nodeStrokeWidth' in network:
node_stroke_width = network['nodeStrokeWidth']
else:
node_stroke_width = cfg['node_stroke_width']
def _set_linewidth_nodes():
glLineWidth(node_stroke_width)
def _set_linewidth_curves():
glLineWidth(cfg['curve_stroke_width'])
def _set_linewidth_legend():
glLineWidth(1.0)
# add the network batch with the right function to set the linewidth
# prior to drawing
window.add_batch(batch, prefunc=_set_linewidth_nodes)
if with_legend:
# add the legend batch with the right function to set the linewidth
# prior to drawing
window.add_batch(legend_batch, prefunc=_set_linewidth_legend)
# decide whether to plot all measured changes or only discrete-time samples
discrete_plot = cfg['plot_sampled_curve']
# find quarantined compartment ids
# This set is needed for filtering later on.
quarantined = set(
model.get_compartment_id(C) for C in quarantine_compartments)
# initialize time arrays
t = 0
discrete_time = [t]
# initialize curves
if with_plot:
# find the maximal value of the
# compartments that are meant to be plotted.
# These sets are needed for filtering later on.
maxy = max([
model.y0[model.get_compartment_id(C)]
for C in (set(model.compartments) - set(ignore_plot_compartments))
])
scl = Scale(bound_increase_factor=cfg['bound_increase_factor'])\
.extent(0,plot_width,plot_height-cfg['padding'],cfg['padding'])\
.domain(0,20*sampling_dt,0,maxy)
curves = {}
for iC, C in enumerate(model.compartments):
if C in ignore_plot_compartments:
continue
_batch = pyglet.graphics.Batch()
window.add_batch(_batch, prefunc=_set_linewidth_curves)
y = [
np.count_nonzero(
model.node_status == model.get_compartment_id(C))
]
curve = Curve(discrete_time, y, cfg['compartment_colors'][iC], scl,
_batch)
curves[C] = curve
# define the pyglet-App update function that's called on every clock cycle
def update(dt):
# skip if nothing remains to be done
if simstate.simulation_ended or simstate.paused:
return
# get sampling_dt
sampling_dt = simstate.sampling_dt
# Advance the simulation until time sampling_dt.
# sim_time is a numpy array including all time values at which
# the system state changed. The first entry is the initial state
# of the simulation at t = 0 which we will discard later on
# the last entry at `sampling_dt` will be missing so we
# have to add it later on.
# `sim_result` is a dictionary that maps a compartment
# to a numpy array containing the compartment counts at
# the corresponding time.
sim_time, sim_result = model.simulate(sampling_dt,
adopt_final_state=True)
# compare the new node statuses with the old node statuses
# and find the nodes that have changed status
ndx = np.where(model.node_status != simstate.old_node_status)[0]
# if nothing changed, evaluate the true total event rate
# and if it's zero, do not do anything anymore
did_simulation_end = len(
ndx) == 0 and model.get_true_total_event_rate() == 0.0
simstate.set_simulation_status(did_simulation_end)
if simstate.simulation_ended:
return
# advance the current time as described above.
# we save both all time values as well as just the sampled times.
this_time = (discrete_time[-1] + sim_time[1:]).tolist() + [
discrete_time[-1] + sampling_dt
]
discrete_time.append(discrete_time[-1] + sampling_dt)
# if curves are plotted
if with_plot:
# iterate through result array
for k, v in sim_result.items():
# skip curves that should be ignored
if k in ignore_plot_compartments:
continue
# count occurrences of this compartment
val = np.count_nonzero(
model.node_status == model.get_compartment_id(k))
if discrete_plot:
# in case only sampled curves are of interest,
# just add this single value
curves[k].append_single_value(discrete_time[-1], v[-1])
else:
# otherwise, append the current value to the exact simulation list
# and append the whole dataset
val = (v[1:].tolist() + [v[-1]])
curves[k].append_list(this_time, val)
# iterate through the nodes that have to be updated
for node in ndx:
status = model.node_status[node]
if cfg['draw_nodes']:
disks[node].color = cfg['compartment_colors'][status]
# if a node becomes quarantined,
# iterate through its attached links (lines)
# and switch them off
if status in quarantined:
for neigh, linkid in node_to_lines[node]:
lines[linkid].visible = False
# if it became unquarantined
elif simstate.old_node_status[node] in quarantined:
# check of the neighbor is unquarantined
# and switch on the link if this is the case
for neigh, linkid in node_to_lines[node]:
if model.node_status[neigh] not in quarantined:
lines[linkid].visible = True
# save the current node statuses
simstate.update(model.node_status)
# schedule the app clock and run the app
pyglet.clock.schedule_interval(update, cfg['update_dt'])
pyglet.app.run()
color=(255, 22, 20),
batch=batch)
rectangle.opacity = 128
rectangle.rotation = 33
line = shapes.Line(100, 100, 100, 200, width=19, batch=batch)
line2 = shapes.Line(150,
150,
444,
111,
width=4,
color=(200, 20, 20),
batch=batch)
arc = shapes.Arc(700,
400,
100,
batch=batch,
closed=True,
angle=2,
color=(255, 0, 0))
print(arc)
arc2 = shapes.Arc(700,
400,
100,
batch=batch,
closed=True,
angle=2,
color=(0, 255, 0))
arc2.rotation = 45
@window.event
def get_network_batch(stylized_network,
yoffset,
draw_links=True,
draw_nodes=True,
draw_nodes_as_rectangles=False,
n_circle_segments=16):
"""
Create a batch for a network visualization.
Parameters
----------
stylized_network : dict
The network properties which are returned from the
interactive visualization.
draw_links : bool, default : True
Whether the links should be drawn
draw_nodes : bool, default : True
Whether the nodes should be drawn
n_circle_segments : bool, default = 16
Number of segments a circle will be constructed of.
Returns
-------
network_objects : dict
A dictionary containing all the necessary objects to draw and
update the network.
- `lines` : a list of pyglet-line objects (one entry per link)
- `disks` : a list of pyglet-circle objects (one entry per node)
- `circles` : a list of pyglet-circle objects (one entry per node)
- `nodes_to_lines` : a dictionary mapping a node to a list of
pairs. Each pair's first entry is the focal node's neighbor
and the second entry is the index of the line-object that
connects the two
- `batch` : the pyglet Batch instance that contains all of the objects
"""
batch = pyglet.graphics.Batch()
pos = {
node['id']: np.array([node['x_canvas'], node['y_canvas'] + yoffset])
for node in stylized_network['nodes']
}
lines = []
disks = []
circles = []
node_to_lines = {node['id']: [] for node in stylized_network['nodes']}
if draw_links:
for ilink, link in enumerate(stylized_network['links']):
u, v = link['source'], link['target']
node_to_lines[u].append((v, ilink))
node_to_lines[v].append((u, ilink))
if 'color' in link.keys():
this_color = link['color']
else:
this_color = stylized_network['linkColor']
lines.append(
shapes.Line(
pos[u][0],
pos[u][1],
pos[v][0],
pos[v][1],
width=link['width'],
color=tuple(bytes.fromhex(this_color[1:])),
batch=batch,
))
lines[-1].opacity = int(255 * stylized_network['linkAlpha'])
if draw_nodes:
disks = [None for n in range(len(stylized_network['nodes']))]
circles = [None for n in range(len(stylized_network['nodes']))]
for node in stylized_network['nodes']:
if not draw_nodes_as_rectangles:
disks[node['id']] = \
shapes.Circle(node['x_canvas'],
node['y_canvas']+yoffset,
node['radius'],
segments=n_circle_segments,
color=tuple(bytes.fromhex(node['color'][1:])),
batch=batch,
)
circles[node['id']] = \
shapes.Arc(node['x_canvas'],
node['y_canvas']+yoffset,
node['radius'],
segments=n_circle_segments+1,
color=tuple(bytes.fromhex(stylized_network['nodeStrokeColor'][1:])),
batch=batch,
)
else:
r = node['radius']
disks[node['id']] = \
shapes.Rectangle(
node['x_canvas']-r,
node['y_canvas']+yoffset-r,
2*r,
2*r,
color=tuple(bytes.fromhex(node['color'][1:])),
batch=batch)
return {
'lines': lines,
'disks': disks,
'circles': circles,
'node_to_lines': node_to_lines,
'batch': batch
}