40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390 | def generate_simulator(args) -> Simulator:
"""Loads the world saved as an hdf5 file, and creates a simulator ready for running.
Originally written in its entirety in run_simulation.py. run_simulation_2.py places it here to make things neater.
Args:
args:
"""
logger.info("Loading the world with world_loader.generate_simulator()")
try:
logger.info("Initialising disease configuration")
# Loading configurations of the disease and other aspects
logger.info("Loading yaml configuration file")
def __load_config(config_path):
"""Load yaml configuration file"""
with open(config_path) as f:
return yaml.load(f, Loader=yaml.FullLoader)
CONFIG_PATH = args.config
config = __load_config(CONFIG_PATH)
disease_settings = config.get("disease")
disease = disease_settings.get("model")
parameters_path = args.parameters
disease_config_path = args.disease_config_path
disease_config = DiseaseConfig(disease, disease_config_path)
GlobalContext.set_disease_config(disease_config)
logger.info("Disease config successfully loaded")
except:
raise Exception("Could not load the disease configuration")
try:
logger.info("Initialising record")
record = Record(
record_path=args.save_path,
record_static_data=True,
mpi_rank=mpi_rank
)
# Domain decomposition is different in MPI vs non-MPI modes
logger.info("Loading world from HDF5")
if mpi_available and mpi_size > 1:
# MPI mode with multiple processes
if mpi_rank == 0:
with h5py.File(args.world_path, "r") as f:
super_area_ids = f["geography"]["super_area_id"][:]
super_area_names = f["geography"]["super_area_name"][:]
super_area_regions = f["geography"]["super_area_region"][:]
region_names = f["geography"]["region_name"][:]
super_area_name_to_id = {
name.decode(): id for name, id in zip(super_area_names, super_area_ids)
}
# Store population data for work-only area detection
n_people = f["geography"]["super_area_n_people"][:]
n_workers = f["geography"]["super_area_n_workers"][:]
# Create super_area_id -> region_name mapping
super_area_ids_to_region = {}
for i, super_area_id in enumerate(super_area_ids):
region_idx = super_area_regions[i]
region_name = region_names[region_idx].decode()
super_area_ids_to_region[int(super_area_id)] = region_name
super_areas_per_domain, score_per_domain = DomainSplitter.generate_world_split(
number_of_domains=mpi_size, world_path=args.world_path
)
super_area_names_to_domain_dict = {}
super_area_ids_to_domain_dict = {}
for domain, super_areas in super_areas_per_domain.items():
for super_area in super_areas:
super_area_names_to_domain_dict[super_area] = domain
# Find ALL IDs for this super area from geography data (handle duplicates)
found_ids = []
for i, geo_name in enumerate(super_area_names):
decoded_name = geo_name.decode() if isinstance(geo_name, bytes) else geo_name
if decoded_name == super_area:
found_ids.append(int(super_area_ids[i]))
if found_ids:
for super_area_id in found_ids:
super_area_ids_to_domain_dict[super_area_id] = domain
if len(found_ids) > 1:
#print(f"DEBUG: Found {len(found_ids)} IDs for {super_area}: {found_ids}")
pass
else:
print(f"WARNING: Could not find ID for super area {super_area}")
# Add work-only super areas (areas with workers but no residents) to domain 0
# These areas are not included in domain decomposition but people work there
work_only_count = 0
# Check ALL areas from geography, not just domain_decomposition_data
for geo_index, (geo_name, geo_id) in enumerate(zip(super_area_names, super_area_ids)):
name = geo_name.decode() if isinstance(geo_name, bytes) else geo_name
id_val = int(geo_id)
people_count = n_people[geo_index]
workers_count = n_workers[geo_index]
# If no residents but has workers, and not already assigned
if people_count == 0 and workers_count > 0 and name not in super_area_names_to_domain_dict:
logger.info(f"Adding work-only super area {name} (ID {id_val}, {workers_count} workers) to domain 0")
super_area_names_to_domain_dict[name] = 0
super_area_ids_to_domain_dict[id_val] = 0
work_only_count += 1
logger.info(f"Added {work_only_count} work-only super areas to domain 0")
with open("super_area_ids_to_domain.json", "w") as f:
json.dump(super_area_ids_to_domain_dict, f)
with open("super_area_names_to_domain.json", "w") as f:
json.dump(super_area_names_to_domain_dict, f)
with open("super_area_ids_to_region.json", "w") as f:
json.dump(super_area_ids_to_region, f)
if mpi_available:
mpi_comm.Barrier()
if mpi_rank > 0:
with open("super_area_ids_to_domain.json", "r") as f:
super_area_ids_to_domain_dict = json.load(f, object_hook=keys_to_int)
# Load super_area -> region mapping on all ranks
with open("super_area_ids_to_region.json", "r") as f:
super_area_ids_to_region_dict = json.load(f, object_hook=keys_to_int)
domain = Domain.from_hdf5(
domain_id=mpi_rank,
super_areas_to_domain_dict=super_area_ids_to_domain_dict,
hdf5_file_path=args.world_path,
interaction_config=args.parameters,
super_areas_to_region_dict=super_area_ids_to_region_dict,
)
else:
# Non-MPI mode or single MPI process - load entire world
logger.info("Loading entire world in a single domain")
# For non-MPI mode, we create a simple domain dict that assigns everything to domain 0
with h5py.File(args.world_path, "r") as f:
super_area_ids = f["geography"]["super_area_id"]
super_area_names = f["geography"]["super_area_name"]
# Create a dictionary mapping all super areas to domain 0
# This includes both residential areas and work-only areas
super_area_ids_to_domain_dict = {int(id): 0 for id in super_area_ids}
logger.info(f"Non-MPI mode: Created domain dictionary with {len(super_area_ids_to_domain_dict)} super areas (includes work-only areas)")
domain = Domain.from_hdf5(
domain_id=0,
super_areas_to_domain_dict=super_area_ids_to_domain_dict,
hdf5_file_path=args.world_path,
interaction_config=args.parameters,
)
logger.info(f"Domain loaded successfully for rank {mpi_rank}")
# Regenerate leisure
logger.info("Generating leisure activities")
try:
leisure = generate_leisure_for_config(domain, CONFIG_PATH)
logger.info("Leisure activities generated successfully")
except Exception as e:
logger.warning(f"Error generating leisure activities: {e}")
logger.warning("Continuing without leisure activities")
leisure = None
# Initialise disease model
logger.info("Setting up disease model and infection selectors")
disease_config = GlobalContext.get_disease_config()
try:
selector = InfectionSelector.from_disease_config(disease_config)
selectors = InfectionSelectors([selector])
logger.info("Infection selectors initialised successfully")
except Exception as e:
logger.error(f"Error initialising infection selectors: {e}")
raise
try:
# Initialise infection seed from configuration
logger.info("Setting up infection seed from configuration")
# Load configuration using the from_file class method
loader = SeedingConfigLoader.from_file(args.seeding_config)
# Validate configuration
errors = loader.validate_config()
if errors:
error_msg = "Configuration validation failed:\n" + "\n".join(errors)
logger.error(error_msg)
raise ValueError(error_msg)
# Create infection seeds
infection_seeds = loader.create_infection_seeds(domain, selector)
logger.info("Infection seeds initialised successfully from configuration")
except Exception as e:
logger.error(f"Error initialising infection seeds: {e}")
raise
# Initialise vaccination campaigns
logger.info("Setting up vaccination campaigns")
try:
vaccination_campaigns = VaccinationCampaigns.from_disease_config(disease_config)
logger.info("Vaccination campaigns initialised successfully")
except Exception as e:
logger.warning(f"Error initialising vaccination campaigns: {e}")
logger.warning("Continuing without vaccination campaigns")
vaccination_campaigns = None
# Initialise epidemiology
logger.info("Setting up epidemiology")
epidemiology = Epidemiology(
infection_selectors=selectors, infection_seeds=infection_seeds, vaccination_campaigns=vaccination_campaigns
)
logger.info("Epidemiology initialised successfully")
# Initialise interaction
logger.info("Setting up interaction model")
try:
interaction = Interaction.from_file()
logger.info("Interaction model initialised successfully")
except Exception as e:
logger.error(f"Error initialising interaction model: {e}")
raise
# Initialise policies
logger.info("Setting up policies")
try:
# Now try to initialise policies
# Load policies using only june.policy
policies = Policies.from_file(
disease_config=disease_config,
base_policy_modules=("june.policy",) # Remove camps.policy
)
logger.info("Policies initialised successfully")
except Exception as e:
logger.error(f"Error initialising policies: {e}")
logger.error("Full traceback:", exc_info=True)
logger.warning("Continuing without policies")
policies = None
# Initialise events
logger.info("Setting up events")
try:
events = Events.from_file()
logger.info("Events initialised successfully")
except Exception as e:
logger.warning(f"Error initialising events: {e}")
logger.warning("Continuing without events")
events = None
# Initialise travel
logger.info("Setting up travel")
try:
travel = Travel()
logger.info("Travel initialised successfully")
except Exception as e:
logger.warning(f"Error initialising travel: {e}")
logger.warning("Continuing without travel")
travel = None
# Check if sexual encounters are enabled in config
sexual_encounter = None
try:
with open(CONFIG_PATH) as f:
config = yaml.safe_load(f)
features = config.get("features", {})
sexual_encounters_config = features.get("sexual_encounters", {"enabled": False})
sexual_encounters_enabled = sexual_encounters_config.get("enabled", False)
if sexual_encounters_enabled:
logger.info("Setting up sexual encounters")
from june.sexual_encounter.sexual_encounter import SexualEncounter
private_rooms = getattr(domain, 'private_rooms', None)
# Pass the CONFIG_PATH to SexualEncounter for configuration loading
sexual_encounter = SexualEncounter(
private_rooms=private_rooms,
contact_manager=None, # Will be set later when simulator creates contact_manager
config_path=None # Uses default path to sexual_relationships_distributor.yaml
)
logger.info("Sexual encounters initialised successfully")
else:
logger.info("Sexual encounters disabled in config")
except Exception as e:
logger.warning(f"Error initialising sexual encounters: {e}")
logger.warning("Continuing without sexual encounters")
sexual_encounter = None
# Inspect domain venues
group_types = []
domainVenues = {}
# Define a mapping of domain attributes and a flag indicating if "bins" is required
venue_attributes = {
"households": True,
"care_homes": True,
"schools": True,
"hospitals": False,
"companies": True,
"universities": True,
"pubs": True,
"groceries": True,
"cinemas": True,
"gyms": True,
"city_transports": False,
"inter_city_transports": False,
}
# Log venue information
logger.info("Inspecting domain venues")
for venue, has_bins in venue_attributes.items():
venue_data = getattr(domain, venue, None)
if venue_data is not None:
if len(venue_data) > 0:
group_types.append(venue_data)
domainVenues[venue] = {
"N": len(venue_data),
"bins": venue_data[0].subgroup_bins if has_bins else None,
}
logger.info(f"Found {len(venue_data)} {venue}")
else:
domainVenues[venue] = {"N": 0, "bins": "NaN" if has_bins else None}
logger.info(f"No {venue} found")
else:
logger.info(f"No {venue} attribute in domain")
logger.info("Domain venue inspection complete")
# Initialise simulator
logger.info("Creating simulator")
simulator = Simulator.from_file(
world=domain,
policies=policies,
events=events,
interaction=interaction,
leisure=leisure,
travel=travel,
sexual_encounter=sexual_encounter,
epidemiology=epidemiology,
config_filename=CONFIG_PATH,
record=record
)
logger.info("Simulator generated successfully")
return simulator
except Exception as e:
logger.error(f"Error generating simulator: {e}")
logger.error(traceback.format_exc())
raise
|