Coverage for tests/models/test_SIR_ABM.py: 100%

1import sys 

2from os.path import dirname as up 

3 

4import pytest 

5import torch 

6from dantro._import_tools import import_module_from_path 

7from dantro._yaml import load_yml 

8from pkg_resources import resource_filename 

9 

10sys.path.insert(0, up(up(up(__file__)))) 

11 

12SIR = import_module_from_path(mod_path=up(up(up(__file__))), mod_str="models.SIR") 

13vec = import_module_from_path(mod_path=up(up(up(__file__))), mod_str="include.vector") 

14 

15Agent = SIR.ABM.Agent 

16Vector = vec.Vector 

17 

18# Load the test config 

19CFG_FILENAME = resource_filename("tests", "cfgs/SIR_Dynamics.yml") 

20test_cfg = load_yml(CFG_FILENAME) 

21 

22 

23def test_agent(): 

24 """Test agent creation and movement""" 

25 

26 agent = Agent(id=0, kind="some_kind", position=Vector(0, 0)) 

27 

28 agent.move(Vector(1, 1)) 

29 assert agent.position == Vector(1, 1) 

30 

31 agent.move(Vector(-1, -1)) 

32 assert agent.position == Vector(0, 0) 

33 

34 space = [[-1, 1], [-1, 1]] 

35 

36 # Test repelling from wall 

37 agent.position = Vector(-0.5, 0) 

38 agent.repel_from_wall(space=space, direction=Vector(-1, 0)) 

39 assert agent.position == Vector(-0.5, 0) 

40 

41 agent.repel_from_wall(space=space, direction=Vector(-1, 1)) 

42 assert agent.position == Vector(-0.5, 1) 

43 

44 # Test movement in periodic space 

45 agent.move_in_periodic_space(direction=Vector(1, 1), space=space) 

46 assert agent.position == Vector(0.5, 0) 

47 

48 agent.move_in_periodic_space(direction=Vector(1, 2), space=space) 

49 assert agent.position == Vector(-0.5, 0) 

50 

51 agent.position = Vector(0, 0) 

52 agent.move_in_periodic_space(direction=Vector(2, 2), space=space) 

53 assert agent.position == Vector(0, 0) 

54 

55 space = Vector(4, 5) 

56 

57 # Test diffusion 

58 agent.move_randomly_in_space(space=space, diffusion_radius=0.5) 

59 assert agent.position.within_space(space) 

60 assert agent.position != Vector(0, 0) 

61 assert abs(agent.position) == pytest.approx(0.5, 1e-10) 

62 

63 agent.move_randomly_in_space(space=space, diffusion_radius=4) 

64 assert agent.position.within_space(space) 

65 agent.position = Vector(0, 0) 

66 agent.move_randomly_in_space(space=space, diffusion_radius=10, periodic=True) 

67 

68 agent.position = Vector(0, 0) 

69 space = [[-1, 1], [-1, 1]] 

70 agent.move_randomly_in_space(space=space, diffusion_radius=1) 

71 assert agent.position != Vector(0, 0) 

72 

73 # Test reset 

74 agent.kind = "some_other_kind" 

75 agent.reset() 

76 assert agent.position == Vector(0, 0) 

77 assert agent.kind == "some_kind" 

78 

79 # Test representation 

80 assert str(agent) == "Agent 0; kind: some_kind; position: (0, 0)" 

81 

82 

83def test_ABM(): 

84 """Test ABM initialisation""" 

85 

86 for entry in test_cfg: 

87 ABM_cfg = test_cfg[entry] 

88 ABM = SIR.SIR_ABM(**ABM_cfg) 

89 

90 assert ABM 

91 assert ABM.N == ABM_cfg["N"] 

92 assert ABM.p_infect == ABM_cfg["p_infect"] 

93 assert ABM.t_infectious == ABM_cfg["t_infectious"] 

94 assert ABM.space == Vector(ABM_cfg["space"][0], ABM_cfg["space"][1]) 

95 

96 assert len(ABM.current_kinds) == ABM.N 

97 assert ( 

98 ABM.current_counts 

99 == torch.tensor([[ABM.N - 1], [1], [0]], dtype=torch.float) 

100 ).all() 

101 

102 # Test the ABM runs and obeys basic properties 

103 for n in range(ABM_cfg["num_steps"]): 

104 ABM.run_single() 

105 

106 # Test the agent count stays constant 

107 assert torch.sum(ABM.current_counts) == ABM.N 

108 assert torch.sum(torch.abs(ABM.current_counts)) == ABM.N 

109 

110 # Test the agents remain within the space 

111 for _, agent in ABM.agents.items(): 

112 assert agent.position.within_space(ABM.space) 

113 

114 

115def test_dynamics(): 

116 """Test basic dynamics work""" 

117 cfg = test_cfg["dynamics"] 

118 ABM = SIR.SIR_ABM(**cfg) 

119 for n in range(cfg["num_steps"]): 

120 ABM.run_single() 

121 assert torch.sum(ABM.current_counts) == ABM.N 

122 

123 assert len(ABM.current_kinds) == ABM.N 

124 assert ( 

125 ABM.current_counts != torch.tensor([[ABM.N - 1], [1], [0]], dtype=torch.float) 

126 ).all() 

127 

128 ABM.reset() 

129 assert ( 

130 ABM.current_counts == torch.tensor([[ABM.N - 1], [1], [0]], dtype=torch.float) 

131 ).all() 

132 

133 

134def test_no_dynamics(): 

135 """Test nothing happens when p_infect is 0 and t_infectious > num_steps""" 

136 

137 cfg = test_cfg["no_dynamics"] 

138 ABM = SIR.SIR_ABM(**cfg) 

139 for n in range(cfg["num_steps"]): 

140 ABM.run_single() 

141 assert torch.sum(ABM.current_counts) == ABM.N 

142 

143 assert len(ABM.current_kinds) == ABM.N 

144 assert ( 

145 ABM.current_counts == torch.tensor([[ABM.N - 1], [1], [0]], dtype=torch.float) 

146 ).all() 

147 

148 ABM.reset() 

149 

150 

151def test_full_recovery(): 

152 """Test all agents make a full recovery""" 

153 

154 cfg = test_cfg["full_recovery"] 

155 ABM = SIR.SIR_ABM(**cfg) 

156 for n in range(cfg["num_steps"]): 

157 ABM.run_single() 

158 

159 assert (ABM.current_counts[1] == 0).all()