/*

* Copyright (C) 2020-2026 MEmilio

*

* Authors: Daniel Abele

*

* Contact: Martin J. Kuehn <[email protected]>

*

* Licensed under the Apache License, Version 2.0 (the "License");

* you may not use this file except in compliance with the License.

* You may obtain a copy of the License at

*

* http://www.apache.org/licenses/LICENSE-2.0

*

* Unless required by applicable law or agreed to in writing, software

* distributed under the License is distributed on an "AS IS" BASIS,

* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

* See the License for the specific language governing permissions and

* limitations under the License.

*/

#include "abm/simulation.h"

#include "benchmark/benchmark.h"

mio::abm::Simulation<> make_simulation(size_t num_persons, std::initializer_list<uint32_t> seeds)

{

auto rng = mio::RandomNumberGenerator();

rng.seed(seeds);

auto model = mio::abm::Model(5);

model.get_rng() = rng;

//create persons at home

const auto mean_home_size = 5.0;

const auto min_home_size = 1;

auto& home_size_distribution = mio::PoissonDistribution<int>::get_instance();

auto home = model.add_location(mio::abm::LocationType::Home);

auto planned_home_size = home_size_distribution(model.get_rng(), mean_home_size);

auto home_size = 0;

for (size_t i = 0; i < num_persons; ++i) {

if (home_size >= std::max(min_home_size, planned_home_size)) {

home = model.add_location(mio::abm::LocationType::Home);

planned_home_size = home_size_distribution(model.get_rng(), mean_home_size);

home_size = 0;

}

auto age = mio::AgeGroup(mio::UniformIntDistribution<size_t>::get_instance()(

model.get_rng(), size_t(0), model.parameters.get_num_groups() - 1));

auto person = model.add_person(home, age);

model.assign_location(uint32_t(i), home);

home_size++;

}

//create other locations

for (auto loc_type :

{mio::abm::LocationType::School, mio::abm::LocationType::Work, mio::abm::LocationType::SocialEvent,

mio::abm::LocationType::BasicsShop, mio::abm::LocationType::Hospital, mio::abm::LocationType::ICU}) {

const auto num_locs = std::max(size_t(1), num_persons / 2'000);

std::vector<mio::abm::LocationId> locs(num_locs);

std::generate(locs.begin(), locs.end(), [&] {

return model.add_location(loc_type);

});

for (size_t p = 0; p < num_persons; ++p) {

auto loc_idx =

mio::UniformIntDistribution<size_t>::get_instance()(model.get_rng(), size_t(0), num_locs - 1);

model.assign_location(uint32_t(p), locs[loc_idx]);

}

}

//infections and masks

for (auto& person : model.get_persons()) {

auto prng = mio::abm::PersonalRandomNumberGenerator(person);

//some % of people are infected, large enough to have some infection activity without everyone dying

auto pct_infected = 0.05;

if (mio::UniformDistribution<ScalarType>::get_instance()(prng, 0.0, 1.0) < pct_infected) {

auto state = mio::abm::InfectionState(

mio::UniformIntDistribution<int>::get_instance()(prng, 1, int(mio::abm::InfectionState::Count) - 1));

auto infection = mio::abm::Infection(prng, mio::abm::VirusVariant::Wildtype, person.get_age(),

model.parameters, mio::abm::TimePoint(0), state);

person.add_new_infection(std::move(infection));

}

//equal chance of (moderate) mask refusal and (moderate) mask eagerness

auto pct_compliance_values = std::array{0.05 /*0*/, 0.2 /*0.25*/, 0.5 /*0.5*/, 0.2 /*0.75*/, 0.05 /*1*/};

auto compliance_value = 0.25 * mio::DiscreteDistribution<int>::get_instance()(prng, pct_compliance_values);

person.set_compliance(mio::abm::InterventionType::Mask, compliance_value);

}

//masks at locations

for (auto& loc : model.get_locations()) {

//some % of locations require masks

//skip homes so persons always have a place to go, simulation might break otherwise

auto pct_require_mask = 0.2;

if (loc.get_type() != mio::abm::LocationType::Home &&

mio::UniformDistribution<ScalarType>::get_instance()(model.get_rng()) < pct_require_mask) {

loc.set_required_mask(mio::abm::MaskType::Community);

}

}

//testing schemes

auto sample = [&](auto v, size_t n) { //selects n elements from list v

std::shuffle(v.begin(), v.end(), model.get_rng());

return std::vector<typename decltype(v)::value_type>(v.begin(), v.begin() + n);

};

std::vector<mio::AgeGroup> ages;

std::generate_n(std::back_inserter(ages), model.parameters.get_num_groups(), [a = 0]() mutable {

return mio::AgeGroup(a++);

});

auto random_criteria = [&]() {

auto random_ages = sample(ages, 2);

auto random_states = std::vector<mio::abm::InfectionState>(0);

return mio::abm::TestingCriteria(random_ages, random_states);

};

model.get_testing_strategy().add_scheme(

{mio::abm::LocationType::School, mio::abm::LocationType::Work, mio::abm::LocationType::SocialEvent,

mio::abm::LocationType::Home},

mio::abm::TestingScheme(random_criteria(), mio::abm::days(3), mio::abm::TimePoint(0),

mio::abm::TimePoint(0) + mio::abm::days(10), {}, 0.5));

return mio::abm::Simulation(mio::abm::TimePoint(0), std::move(model));

}

/**

* Benchmark for the ABM simulation.

* @param num_persons Number of persons in the simulation.

* @param seeds Seeds for the random number generator.

*/

void abm_benchmark(benchmark::State& state, size_t num_persons, std::initializer_list<uint32_t> seeds)

{

mio::set_log_level(mio::LogLevel::warn);

for (auto&& _ : state) {

state.PauseTiming(); //exclude the setup from the benchmark

auto sim = make_simulation(num_persons, seeds);

state.ResumeTiming();

//simulated time should be long enough to have full infection runs and mobility to every location

auto final_time = sim.get_time() + mio::abm::days(10);

sim.advance(final_time);

//debug output can be enabled to check for unexpected results (e.g. infections dieing out)

//normally should have no significant effect on runtime

const bool monitor_infection_activity = false;

if constexpr (monitor_infection_activity) {

std::cout << "num_persons = " << num_persons << "\n";

for (auto inf_state = 0; inf_state < (int)mio::abm::InfectionState::Count; inf_state++) {

std::cout << "inf_state = " << inf_state << ", sum = "

<< sim.get_model().get_subpopulation_combined(sim.get_time(),

mio::abm::InfectionState(inf_state))

<< "\n";

}

}

}

}

//Measure ABM simulation run time with different sizes and different seeds.

//Fixed RNG seeds to make runs comparable. When there are code changes, the simulation will still

//run differently due to different sequence of random numbers being drawn. But for large enough sizes

//RNG should average out, so runs should be comparable even with code changes.

//We run a few different benchmarks to hopefully catch abnormal cases. Then seeds may

//have to be adjusted to get the benchmark back to normal.

//For small sizes (e.g. 10k) extreme cases are too likely, i.e. infections die out

//or overwhelm everything, so we don't benchmark these. Results should be mostly transferrable.

BENCHMARK_CAPTURE(abm_benchmark, abm_benchmark_50k, 50000, {14159265u, 35897932u})->Unit(benchmark::kMillisecond);

BENCHMARK_CAPTURE(abm_benchmark, abm_benchmark_100k, 100000, {38462643u, 38327950u})->Unit(benchmark::kMillisecond);

BENCHMARK_CAPTURE(abm_benchmark, abm_benchmark_200k, 200000, {28841971u, 69399375u})->Unit(benchmark::kMillisecond);

BENCHMARK_MAIN();