def test_go_record():
bristol = Ward("bristol")
london = Ward("london")
oxford = Ward("oxford")
bristol.set_num_players(100)
london.set_num_players(100)
oxford.set_num_players(100)
bristol.add_workers(0, destination=london)
bristol.add_workers(0, destination=oxford)
oxford.add_workers(0, destination=bristol)
oxford.add_workers(0, destination=london)
london.add_workers(0, destination=bristol)
london.add_workers(0, destination=oxford)
disease = Disease(name="lurgy")
disease.add(name="E", beta=0.5, progress=0.5)
disease.add(name="I", beta=0.8, progress=0.25)
disease.add(name="R")
disease.assert_sane()
params = Parameters()
params.set_disease(disease)
params.add_seeds("1 20 bristol")
outdir = os.path.join(script_dir, "test_go_ward_output")
network = Network.from_wards(bristol + london + oxford, params=params)
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.copy().run(population=Population(),
output_dir=output_dir,
mover=move_bristol)
OutputFiles.remove(outdir, prompt=None)
# despite the moves, only Bristolians can be infected as there
# are no work movements
assert trajectory[-1].recovereds <= 100
def test_bg_foi():
bristol = Ward("Bristol")
london = Ward("London")
oxford = Ward("Oxford")
bristol.set_num_players(100)
london.set_num_players(100)
oxford.set_num_players(100)
bristol.set_bg_foi(10)
london.set_bg_foi(1)
wards = bristol + london + oxford
lurgy = Disease("lurgy")
lurgy.add("E", beta=0.0, progress=1.0)
lurgy.add("I", beta=0.0, progress=1.0)
lurgy.add("R")
params = Parameters()
params.set_disease(lurgy)
network = Network.from_wards(wards, params=params)
outdir = os.path.join(script_dir, "test_bg_foi")
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.run(population=Population(),
output_dir=output_dir,
extractor=_check_infections)
OutputFiles.remove(outdir, prompt=None)
# the only source of infection should be the bg. Should have
# no infections in oxford, some in london and most in bristol
print(trajectory[-1])
assert trajectory[-1].recovereds <= 200
assert trajectory[-1].susceptibles >= 100
def test_go_ward():
bristol = Ward("bristol")
london = Ward("london")
oxford = Ward("oxford")
bristol.set_num_players(100)
london.set_num_players(100)
oxford.set_num_players(100)
bristol.add_workers(0, destination=london)
bristol.add_workers(0, destination=oxford)
oxford.add_workers(0, destination=bristol)
oxford.add_workers(0, destination=london)
london.add_workers(0, destination=bristol)
london.add_workers(0, destination=oxford)
disease = Disease(name="lurgy")
disease.add(name="E", beta=0.5, progress=0.5)
disease.add(name="I", beta=0.8, progress=0.25)
disease.add(name="R")
disease.assert_sane()
params = Parameters()
params.set_disease(disease)
params.add_seeds("1 20 bristol")
outdir = os.path.join(script_dir, "test_go_ward_output")
network = Network.from_wards(bristol + london + oxford, params=params)
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.copy().run(population=Population(),
output_dir=output_dir,
nthreads=1)
print(trajectory[-1])
# there are no workers, so can only infect the 100 Bristolians
assert trajectory[-1].recovereds <= 100
# Repeat with a different number of threads to test parallel code
outdir = os.path.join(script_dir, "test_go_ward_output")
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.copy().run(population=Population(),
output_dir=output_dir,
nthreads=4)
print(trajectory[-1])
# there are no workers, so can only infect the 100 Bristolians
assert trajectory[-1].recovereds <= 100
# now do a cyclic move, from "R" to "S"
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.copy().run(population=Population(),
output_dir=output_dir,
mover=move_cycle,
nsteps=200)
print(trajectory[-1])
# we should have completed all 200 steps as this will never end
assert trajectory[-1].day == 200
assert trajectory[-1].recovereds <= 100
# move everyone to Bristol
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.copy().run(population=Population(),
output_dir=output_dir,
mover=move_bristol)
print(trajectory[-1])
assert trajectory[-1].recovereds > 100
# now move through all of bristol, london, oxford
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.copy().run(population=Population(),
output_dir=output_dir,
mover=move_travel)
print(trajectory[-1])
assert trajectory[-1].recovereds > 100
OutputFiles.remove(outdir, prompt=None)
def test_ward_conversion():
# load all of the parameters
try:
params = Parameters.load(parameters="march29")
except Exception as e:
print(f"Unable to load parameter files. Make sure that you have "
f"cloned the MetaWardsData repository and have set the "
f"environment variable METAWARDSDATA to point to the "
f"local directory containing the repository, e.g. the "
f"default is $HOME/GitHub/MetaWardsData")
raise e
params.set_input_files("2011Data")
print("Building the network...")
network = Network.build(params=params)
profiler = Profiler()
profiler = profiler.start("to_json")
wards = network.to_wards(profiler=profiler)
print(f"{wards.num_workers()} / {wards.num_players()}")
_assert_equal(wards.num_workers(), network.work_population)
_assert_equal(wards.num_players(), network.play_population)
print(f"{wards.num_work_links()} / {wards.num_play_links()}")
_assert_equal(wards.num_work_links(), network.nlinks)
_assert_equal(wards.num_play_links(), network.nplay)
print("Converting to data...")
data = wards.to_data(profiler=profiler)
print("Converting to json...")
profiler = profiler.start("Convert to JSON")
s = json.dumps(data)
profiler = profiler.stop()
profiler = profiler.stop() # end to_json
print(f"Done - {len(s)/(1024*1024.0)} MB : {s[0:1024]}...")
print(f"Converting from json...")
profiler = profiler.start("Convert from JSON")
profiler = profiler.start("from_json")
data = json.loads(s)
profiler = profiler.stop()
wards2 = Wards.from_data(data, profiler=profiler)
assert wards2 == wards
network2 = Network.from_wards(wards2, profiler=profiler)
profiler = profiler.stop()
Console.print(profiler)
Console.print("Validating equality - may take some time...")
_assert_equal(network2.nnodes, network.nnodes)
_assert_equal(network2.nlinks, network.nlinks)
_assert_equal(network2.nplay, network.nplay)
if network.info is None:
assert network2.info is None
_assert_equal(len(network.info), len(network2.info))
Console.print(f"{len(network.info)}, {network.nnodes}")
with Console.progress() as progress:
task1 = progress.add_task("Validating info", total=len(network.info))
task2 = progress.add_task("Validating nodes", total=network.nnodes)
task3 = progress.add_task("Validating work", total=network.nlinks)
task4 = progress.add_task("Validating play", total=network.nplay)
for i in range(0, len(network.info)):
assert network.info[i] == network2.info[i]
progress.update(task1, advance=1)
progress.update(task1, completed=len(network.info), force_update=True)
for i in range(1, network.nnodes + 1):
_assert_equal(network.nodes.label[i], network2.nodes.label[i])
_assert_equal(network.nodes.begin_to[i],
network2.nodes.begin_to[i])
_assert_equal(network.nodes.end_to[i], network2.nodes.end_to[i])
_assert_equal(network.nodes.self_w[i], network2.nodes.self_w[i])
_assert_equal(network.nodes.begin_p[i], network2.nodes.begin_p[i])
_assert_equal(network.nodes.end_p[i], network2.nodes.end_p[i])
_assert_equal(network.nodes.self_p[i], network2.nodes.self_p[i])
_assert_equal(network.nodes.x[i], network2.nodes.x[i])
_assert_equal(network.nodes.y[i], network2.nodes.y[i])
progress.update(task2, advance=1)
progress.update(task2, completed=network.nnodes, force_update=True)
for i in range(1, network.nlinks + 1):
_assert_equal(network.links.ifrom[i], network2.links.ifrom[i])
_assert_equal(network.links.ito[i], network2.links.ito[i])
_assert_equal(network.links.weight[i], network2.links.weight[i])
_assert_equal(network.links.suscept[i], network2.links.suscept[i])
progress.update(task3, advance=1)
progress.update(task3, completed=network.nlinks, force_update=True)
for i in range(1, network.nplay + 1):
_assert_equal(network.play.ifrom[i], network2.play.ifrom[i])
_assert_equal(network.play.ito[i], network2.play.ito[i])
_assert_equal(network.play.weight[i], network2.play.weight[i])
_assert_equal(network.play.suscept[i], network2.play.suscept[i])
progress.update(task4, advance=1)
progress.update(task4, completed=network.nplay, force_update=True)
def test_ward_params():
bristol = Ward("Bristol")
london = Ward("London")
oxford = Ward("Oxford")
bristol.set_num_players(1000)
london.set_num_players(5000)
oxford.set_num_players(800)
bristol.add_workers(500, destination=bristol)
bristol.add_workers(200, destination=oxford)
bristol.add_workers(900, destination=london)
oxford.add_workers(100, destination=bristol)
oxford.add_workers(1000, destination=oxford)
oxford.add_workers(2000, destination=london)
london.add_workers(100, destination=bristol)
london.add_workers(400, destination=oxford)
london.add_workers(10000, destination=london)
bristol.add_player_weight(0.1, destination=oxford)
oxford.add_player_weight(0.2, destination=bristol)
oxford.add_player_weight(0.15, destination=london)
london.add_player_weight(0.3, destination=oxford)
bristol.set_scale_uv(0.5)
bristol.set_cutoff(5.0)
london.set_scale_uv(1.5)
oxford.set_cutoff(15.4)
bristol.set_bg_foi(10.8)
london.set_bg_foi(-5.2)
bristol.set_custom("something", 9.9)
oxford.set_custom("else", 22.1)
assert bristol.scale_uv() == 0.5
assert bristol.cutoff() == 5.0
assert bristol.bg_foi() == 10.8
assert bristol.custom("something") == 9.9
assert bristol.custom("else") == 0.0
assert bristol.custom("else", 1.0) == 1.0
assert london.scale_uv() == 1.5
assert london.bg_foi() == -5.2
assert oxford.cutoff() == 15.4
assert oxford.bg_foi() == 0.0
assert oxford.custom("else") == 22.1
wards = bristol + london + oxford
assert wards[bristol].scale_uv() == 0.5
assert wards[oxford].custom("else") == 22.1
assert wards[bristol].bg_foi() == 10.8
assert wards[oxford].bg_foi() == 0.0
s = wards.to_json()
print(s)
assert wards == Wards.from_json(s)
network = Network.from_wards(wards)
assert network.nodes.scale_uv[1] == 0.5
assert network.nodes.scale_uv[2] == 1.5
assert network.nodes.scale_uv[3] == 1.0
assert network.nodes.cutoff[1] == 5.0
assert network.nodes.cutoff[2] == 99999.99
assert network.nodes.cutoff[3] == 15.4
assert network.nodes.bg_foi[1] == 10.8
assert network.nodes.bg_foi[2] == -5.2
assert network.nodes.bg_foi[3] == 0.0
something = network.nodes.get_custom("something")
els = network.nodes.get_custom("else")
miss = network.nodes.get_custom("missing", 99.9)
assert something[1] == 9.9
assert something[2] == 0.0
assert something[3] == 0.0
assert els[1] == 0.0
assert els[2] == 0.0
assert els[3] == 22.1
assert miss[1] == 99.9
assert miss[2] == 99.9
assert miss[3] == 99.9
network2 = Network.from_wards(network.to_wards())
for i in range(0, len(network.info)):
assert network.info[i] == network2.info[i]
for i in range(1, network.nnodes + 1):
assert network.nodes.scale_uv[i] == network2.nodes.scale_uv[i]
assert network.nodes.cutoff[i] == network2.nodes.cutoff[i]
assert network.nodes.bg_foi[i] == network2.nodes.bg_foi[i]
assert network.nodes.get_custom("something")[i] == \
network2.nodes.get_custom("something")[i]
assert network.nodes.get_custom("else")[i] == \
network2.nodes.get_custom("else")[i]
def test_wards():
wards = Wards()
bristol = Ward(name="Bristol")
london = Ward(name="London")
bristol.add_workers(500, destination=bristol)
bristol.add_workers(500, destination=london)
bristol.set_num_players(750)
london.add_workers(8500, destination=london)
london.add_workers(100, destination=bristol)
london.set_num_players(10000)
print(bristol)
print(london)
wards.add(bristol)
wards.add(london)
print(wards)
assert wards.num_workers() == 9600
assert wards.num_players() == 10750
assert wards.population() == 10750 + 9600
data = wards.to_data()
print(data)
wards2 = Wards.from_data(data)
print(wards2)
assert wards == wards2
try:
filename = "_test_wards.json"
with open(filename, "w") as FILE:
json.dump(data, FILE, indent=2)
print("".join(open(filename).readlines()))
with open(filename) as FILE:
data2 = json.load(FILE)
wards2 = Wards.from_data(data2)
print(wards2)
assert wards == wards2
finally:
import os
os.unlink(filename)
network = Network.from_wards(wards, disease="lurgy")
assert network.work_population == wards.num_workers()
assert network.play_population == wards.num_players()
assert network.population == wards.population()
wards2 = network.to_wards()
assert len(wards) == len(wards2)
for w1, w2 in zip(wards, wards2):
if w1 is None:
assert w2 is None
continue
assert w1.id() == w2.id()
assert w1.name() == w2.name()
assert w1.info() == w2.info()
assert w1.num_workers() == w2.num_workers()
assert w1.num_players() == w2.num_players()
assert w1.get_player_lists() == w2.get_player_lists()
assert w1.get_worker_lists() == w2.get_worker_lists()
assert w1.position() == w2.position()
s = wards.to_json()
wards2 = Wards.from_json(s)
assert wards2 == wards
filename = wards.to_json("_test_wards_data.json", indent=4)
print(filename)
try:
wards2 = Wards.from_json(filename)
finally:
import os
os.unlink(filename)
assert wards2 == wards
def test_move_generator():
m = MoveGenerator(fraction=0.3)
assert m.fraction() == 0.3
m = MoveGenerator(number=1000)
assert m.number() == 1000
params = Parameters()
params.set_disease("lurgy")
network = Network.single(params, Population(initial=1000))
m = MoveGenerator(from_stage="E")
assert m.generate(network) == [[0, 1, 0, 1]]
assert m.should_move_all()
assert m.fraction() == 1.0
assert m.number() >= 1000000000
m = MoveGenerator(from_stage="E", to_stage="R")
assert m.generate(network) == [[0, 1, 0, 4]]
assert m.should_move_all()
m = MoveGenerator(to_stage="R")
assert m.generate(network) == [[0, -1, 0, 4], [0, 0, 0, 4], [0, 1, 0, 4],
[0, 2, 0, 4], [0, 3, 0, 4], [0, 4, 0, 4]]
assert m.should_move_all()
m = MoveGenerator(from_stage=["E", "I1"])
assert m.generate(network) == [[0, 1, 0, 1], [0, 2, 0, 2]]
assert m.should_move_all()
m = MoveGenerator(from_stage=["E", "I1"], to_stage=["I1", "I2"])
assert m.generate(network) == [[0, 1, 0, 2], [0, 2, 0, 3]]
assert m.should_move_all()
m = MoveGenerator(from_stage=["E", "I1"], to_stage="R")
assert m.generate(network) == [[0, 1, 0, 4], [0, 2, 0, 4]]
assert m.should_move_all()
m = MoveGenerator()
assert m.generate(network) == [[0, -1, 0, -1], [0, 0, 0, 0], [0, 1, 0, 1],
[0, 2, 0, 2], [0, 3, 0, 3], [0, 4, 0, 4]]
assert m.should_move_all()
demographics = Demographics.load(demography_json)
networks = demographics.specialise(network)
m = MoveGenerator()
assert m.generate(networks) == [[0, -1, 0, -1], [0, 0, 0, 0], [0, 1, 0, 1],
[0, 2, 0, 2], [0, 3, 0, 3], [0, 4, 0, 4],
[1, -1, 1, -1], [1, 0, 1, 0], [1, 1, 1, 1],
[1, 2, 1, 2], [1, 3, 1, 3], [1, 4, 1, 4],
[2, -1, 2, -1], [2, 0, 2, 0], [2, 1, 2, 1],
[2, 2, 2, 2], [2, 3, 2, 3], [2, 4, 2, 4]]
m = MoveGenerator(from_stage="E")
assert m.generate(networks) == [[0, 1, 0, 1], [1, 1, 1, 1], [2, 1, 2, 1]]
m = MoveGenerator(to_stage="R")
assert m.generate(networks) == [[0, -1, 0, 4], [0, 0, 0, 4], [0, 1, 0, 4],
[0, 2, 0, 4], [0, 3, 0, 4], [0, 4, 0, 4],
[1, -1, 1, 4], [1, 0, 1, 4], [1, 1, 1, 4],
[1, 2, 1, 4], [1, 3, 1, 4], [1, 4, 1, 4],
[2, -1, 2, 4], [2, 0, 2, 4], [2, 1, 2, 4],
[2, 2, 2, 4], [2, 3, 2, 4], [2, 4, 2, 4]]
m = MoveGenerator(from_demographic="red")
assert m.generate(networks) == [[0, -1, 0, -1], [0, 0, 0, 0], [0, 1, 0, 1],
[0, 2, 0, 2], [0, 3, 0, 3], [0, 4, 0, 4]]
m = MoveGenerator(to_demographic="blue")
assert m.generate(networks) == [[0, -1, 2, -1], [0, 0, 2, 0], [0, 1, 2, 1],
[0, 2, 2, 2], [0, 3, 2, 3], [0, 4, 2, 4],
[1, -1, 2, -1], [1, 0, 2, 0], [1, 1, 2, 1],
[1, 2, 2, 2], [1, 3, 2, 3], [1, 4, 2, 4],
[2, -1, 2, -1], [2, 0, 2, 0], [2, 1, 2, 1],
[2, 2, 2, 2], [2, 3, 2, 3], [2, 4, 2, 4]]
m = MoveGenerator(to_demographic="blue", to_stage="R")
assert m.generate(networks) == [[0, -1, 2, 4], [0, 0, 2, 4], [0, 1, 2, 4],
[0, 2, 2, 4], [0, 3, 2, 4], [0, 4, 2, 4],
[1, -1, 2, 4], [1, 0, 2, 4], [1, 1, 2, 4],
[1, 2, 2, 4], [1, 3, 2, 4], [1, 4, 2, 4],
[2, -1, 2, 4], [2, 0, 2, 4], [2, 1, 2, 4],
[2, 2, 2, 4], [2, 3, 2, 4], [2, 4, 2, 4]]
m = MoveGenerator(from_demographic="red",
to_demographic="blue",
from_stage=["I1", "I2"],
to_stage="S")
assert m.generate(networks) == [[0, 2, 2, -1], [0, 3, 2, -1]]
m = MoveGenerator(from_demographic="blue", from_stage="E", to_stage="I1")
print(m.generate(networks))
assert m.generate(networks) == [[2, 1, 2, 2]]
m = MoveGenerator(to_ward=1)
player = PersonType.PLAYER
worker = PersonType.WORKER
assert m.generate_wards(network) == [[None, (player, 1, 2)]]
assert m.generate_wards(networks) == [[None, (player, 1, 2)]]
assert not m.should_move_all()
assert m.fraction() == 1.0
assert m.number() >= 1000000000
bristol = Ward("bristol")
london = Ward("london")
oxford = Ward("oxford")
bristol.add_workers(0, destination=london)
bristol.add_workers(0, destination=oxford)
london.add_workers(0, destination=bristol)
london.add_workers(0, destination=oxford)
network = Network.from_wards(bristol + london + oxford)
assert len(network.nodes) == network.nnodes + 1
assert len(network.links) == network.nlinks + 1
assert network.nnodes == 3
assert network.nlinks == 4
m = MoveGenerator(from_ward=["bristol",
WardID("bristol", "london")],
to_ward="oxford",
fraction=0.3,
number=100)
assert not m.should_move_all()
assert m.fraction() == 0.3
assert m.number() == 100
print(m.generate_wards(network))
assert m.generate_wards(network) == [[(player, 1, 2), (player, 3, 4)],
[(worker, 1, 2), (player, 3, 4)]]
m = MoveGenerator(from_ward=WardID("oxford", "bristol"), to_ward=1)
assert m.generate_wards(network) == [[(worker, -1, -1), (player, 1, 2)]]
oxford.add_workers(0, destination=bristol)
network = Network.from_wards(bristol + london + oxford)
print(m.generate_wards(network))
assert m.generate_wards(network) == [[(worker, 5, 6), (player, 1, 2)]]
m = MoveGenerator(from_ward=WardID("bristol", "bristol"), to_ward="oxford")
print(m.generate_wards(network))
assert m.generate_wards(network) == [[(worker, -1, -1), (player, 3, 4)]]
bristol.add_workers(0, destination=bristol)
network = Network.from_wards(bristol + london + oxford)
print(m.generate_wards(network))
assert m.generate_wards(network) == [[(worker, 1, 2), (player, 3, 4)]]
assert network.links.ifrom[1] == network.links.ito[1]
assert network.links.ifrom[1] == network.get_node_index("bristol")
m = MoveGenerator(to_ward=WardID("london", "bristol"))
print(m.generate_wards(network))
assert m.generate_wards(network) == [[None, (worker, 4, 5)]]
assert network.links.ifrom[4] == network.get_node_index("london")
assert network.links.ito[4] == network.get_node_index("bristol")
m = MoveGenerator(
from_stage="S",
to_stage="R",
from_ward=[WardID("bristol", all_commute=True),
WardID("bristol")],
to_ward="oxford",
number=42,
fraction=0.5)
print(m.generate(network))
print(m.generate_wards(network))
assert m.generate(network) == [[0, -1, 0, 4]]
assert m.generate_wards(network) == [[(worker, 1, 4), (player, 3, 4)],
[(player, 1, 2), (player, 3, 4)]]
record = MoveRecord()
for stage in m.generate(network):
for ward in m.generate_wards(network):
from_type = ward[0][0]
ifrom_begin = ward[0][1]
ifrom_end = ward[0][2]
to_type = ward[1][0]
ito_begin = ward[1][1]
ito_end = ward[1][2]
n = ifrom_end - ifrom_begin
print(ifrom_begin, ifrom_end, ito_begin, ito_end)
if ito_end - ito_begin == 0:
raise ValueError("Cannot move individuals to a non-existent "
"ward or ward-link")
elif ito_end - ito_begin == 1:
# this is a single to-ward (or link)
ito_delta = 0
elif ito_end - ito_begin != ifrom_end - ifrom_begin:
# different number of links
raise ValueError(
"Cannot move individuals as the number of from "
"and to links are not the same: "
f"{ifrom_begin}:{ifrom_end} versus "
f"{ito_begin}:{ito_end}")
else:
ito_delta = 1
for i in range(0, n):
ifrom = ifrom_begin + i
ito = ito_begin + (i * ito_delta)
record.add(from_demographic=stage[0],
from_stage=stage[1],
to_demographic=stage[2],
to_stage=stage[3],
from_type=from_type,
from_ward=ifrom,
to_type=to_type,
to_ward=ito,
number=m.number() * m.fraction())
from array import array
records = [
array("i", (0, -1, 1, 1, 0, 4, 2, 3, 21)),
array("i", (0, -1, 1, 2, 0, 4, 2, 3, 21)),
array("i", (0, -1, 1, 3, 0, 4, 2, 3, 21)),
array("i", (0, -1, 2, 1, 0, 4, 2, 3, 21))
]
print(record._record)
for i, move in enumerate(record):
print(f"{move} == {records[i]}")
assert move == records[i]