def test_plugins():
iterator = f"{script_dir}/plugins/_iterator.py"
extractor = f"{script_dir}/plugins/_extractor.py"
mover = f"{script_dir}/plugins/_mover.py"
mixer = f"{script_dir}/plugins/_mixer.py"
params = Parameters()
params.set_disease("ncov")
population = Population(initial=1000)
print("Building the network...")
network = Network.single(params=params, population=population)
outdir = os.path.join(script_dir, "test_plugins")
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.copy().run(population=population,
output_dir=output_dir,
iterator=iterator,
extractor=extractor,
mixer=mixer,
mover=mover,
nthreads=1)
OutputFiles.remove(outdir, prompt=None)
p = trajectory[-1]
assert p.advance_test == "advance_test_value"
assert p.output_test == "output_test_value"
assert p.merge_test == "merge_test_value"
assert p.go_test == "go_test_value"
def test_variableset():
v1 = VariableSet()
v2 = VariableSet()
assert len(v1) == 0
assert len(v2) == 0
with pytest.raises(KeyError):
v1["beta[1]"]
v1["beta[1]"] = 0.5
assert v1["beta[1]"] == 0.5
assert len(v1) == 1
assert v1 != v2
v2["beta[1]"] = 0.5
assert v1 == v2
d = Disease.load("ncov")
p = Parameters()
p.set_disease("ncov")
assert p.disease_params == d
assert p.disease_params.beta[1] != 0.5
p = p.set_variables(v1)
assert p.disease_params.beta[1] == 0.5
def test_set_variables():
d = Disease.load("ncov")
p = Parameters()
p.set_disease("ncov")
assert p.disease_params == d
variables = VariableSet(l1)
p2 = p.set_variables(variables)
assert p.disease_params == d
assert p2.disease_params != d
assert p2.disease_params.beta[2] == 0.9
assert p2.disease_params.beta[3] == 0.93
assert p2.disease_params.progress[1] == 0.18
assert p2.disease_params.progress[2] == 0.92
assert p2.disease_params.progress[3] == 0.90
variables = VariableSet(l0)
p3 = p2.set_variables(variables)
assert p3.disease_params.beta[2] == 0.95
assert p3.disease_params.beta[3] == 0.95
assert p3.disease_params.progress[1] == 0.19
assert p3.disease_params.progress[2] == 0.91
assert p3.disease_params.progress[3] == 0.91
def shared_network():
from metawards import Network, Parameters
params = Parameters()
params.set_input_files("2011Data")
params.set_disease("ncov")
network = Network.build(params, profiler=None)
return network
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_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_not_infected():
# create a disease where it will finish with everyone in the V state
lurgy = Disease("lurgy")
lurgy.add("E", beta=0.0, progress=1.0)
lurgy.add("I1", beta=0.4, progress=0.2)
lurgy.add("I2", beta=0.5, progress=0.5, too_ill_to_move=0.5)
lurgy.add("I3", beta=0.5, progress=0.8, too_ill_to_move=0.8)
lurgy.add("V", beta=0.0, progress=0.0, is_infected=False)
lurgy.add("R")
params = Parameters()
params.set_input_files("single")
params.add_seeds("5")
params.set_disease(lurgy)
network = Network.build(params=params, population=Population(1000))
outdir = os.path.join(script_dir, "test_not_infected")
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.copy().run(population=Population(),
output_dir=output_dir,
nsteps=1000,
nthreads=1)
# should finish at about 60 days and not run to 1000
p = trajectory[-1]
assert p.day < 900
# there should be no recovereds, as all who were ill went to 'others'
assert p.recovereds == 0
assert p.others == p.population - p.susceptibles
# Now repeat with more threads
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.copy().run(population=Population(),
output_dir=output_dir,
nsteps=1000,
nthreads=8)
# should finish at about 60 days and not run to 1000
p = trajectory[-1]
assert p.day < 900
# there should be no recovereds, as all who were ill went to 'others'
assert p.recovereds == 0
assert p.others == p.population - p.susceptibles
# now repeat with demographics
demographics = Demographics.load(redblue_json)
network.params.add_seeds("demographic,number\\n1,10")
network = network.specialise(demographics)
with OutputFiles(outdir, force_empty=True, prompt=None) as output_dir:
trajectory = network.copy().run(population=Population(),
output_dir=output_dir,
nsteps=1000)
# should finish at about 60 days and not run to 1000
p = trajectory[-1]
assert p.day < 900
assert p.recovereds == 0
assert p.others > 0
assert p.others == p.population - p.susceptibles
p_red = p.subpops[0]
p_blue = p.subpops[1]
assert p_red.day < 900
assert p_red.recovereds == 0
assert p_red.others > 0
assert p_red.others == p_red.population - p_red.susceptibles
assert p_blue.day < 900
assert p_blue.recovereds == 0
assert p_blue.others > 0
assert p_blue.others == p_blue.population - p_blue.susceptibles
OutputFiles.remove(outdir, prompt=None)
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_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]