Example #1
0
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"
Example #2
0
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
Example #3
0
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
Example #4
0
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
Example #5
0
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
Example #6
0
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
Example #7
0
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)
Example #8
0
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)
Example #9
0
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]