756 ABM benchmark (#828)

dabele · web-flow · commit 9142b19239a9 · 2023-11-22T17:14:56.000+01:00
Add ABM benchmark, needs to be added to CI later.
diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt
@@ -50,3 +50,6 @@ add_executable(simulation_benchmark simulation.cpp)
 target_link_libraries(simulation_benchmark PRIVATE memilio ode_secir benchmark::benchmark)
 add_executable(graph_simulation_benchmark graph_simulation.cpp)
 target_link_libraries(graph_simulation_benchmark PRIVATE memilio ode_secirvvs benchmark::benchmark)
+
+add_executable(abm_benchmark abm.cpp)
+target_link_libraries(abm_benchmark PRIVATE abm benchmark::benchmark)
diff --git a/cpp/benchmarks/abm.cpp b/cpp/benchmarks/abm.cpp
@@ -0,0 +1,157 @@
+#include "abm/simulation.h"
+#include "memilio/utils/stl_util.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 world      = mio::abm::World(5);
+    world.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                    = world.add_location(mio::abm::LocationType::Home);
+    auto planned_home_size       = home_size_distribution(world.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              = world.add_location(mio::abm::LocationType::Home);
+            planned_home_size = home_size_distribution(world.get_rng(), mean_home_size);
+            home_size         = 0;
+        }
+
+        auto age     = mio::AgeGroup(mio::UniformIntDistribution<size_t>::get_instance()(
+            world.get_rng(), size_t(0), world.parameters.get_num_groups() - 1));
+        auto& person = world.add_person(home, age);
+        person.set_assigned_location(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 world.add_location(loc_type);
+        });
+        for (auto& person : world.get_persons()) {
+            auto loc_idx =
+                mio::UniformIntDistribution<size_t>::get_instance()(world.get_rng(), size_t(0), num_locs - 1);
+            person.set_assigned_location(locs[loc_idx]);
+        }
+    }
+
+    //infections and masks
+    for (auto& person : world.get_persons()) {
+        auto prng = mio::abm::Person::RandomNumberGenerator(world.get_rng(), person);
+        //some % of people are infected, large enough to have some infection activity without everyone dying
+        auto pct_infected = 0.05;
+        if (mio::UniformDistribution<double>::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(),
+                                    world.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_mask_values = std::array{0.05 /*-1*/, 0.2 /*-0.5*/, 0.5 /*0*/, 0.2 /*0.5*/, 0.05 /*1*/};
+        auto mask_value      = -1 + 0.5 * mio::DiscreteDistribution<int>::get_instance()(prng, pct_mask_values);
+        person.set_mask_preferences({size_t(mio::abm::LocationType::Count), mask_value});
+    }
+
+    //masks at locations
+    for (auto& loc : world.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;
+        auto requires_mask    = loc.get_type() != mio::abm::LocationType::Home &&
+                             mio::UniformDistribution<double>::get_instance()(world.get_rng()) < pct_require_mask;
+        loc.set_npi_active(requires_mask);
+    }
+
+    //testing schemes
+    auto sample = [&](auto v, size_t n) { //selects n elements from list v
+        std::shuffle(v.begin(), v.end(), world.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), world.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);
+    };
+
+    world.get_testing_strategy().add_testing_scheme(
+        mio::abm::LocationType::School,
+        mio::abm::TestingScheme(random_criteria(), mio::abm::days(3), mio::abm::TimePoint(0),
+                                mio::abm::TimePoint(0) + mio::abm::days(10), {}, 0.5));
+    world.get_testing_strategy().add_testing_scheme(
+        mio::abm::LocationType::Work,
+        mio::abm::TestingScheme(random_criteria(), mio::abm::days(3), mio::abm::TimePoint(0),
+                                mio::abm::TimePoint(0) + mio::abm::days(10), {}, 0.5));
+    world.get_testing_strategy().add_testing_scheme(
+        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));
+    world.get_testing_strategy().add_testing_scheme(
+        mio::abm::LocationType::SocialEvent,
+        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(world));
+}
+
+/**
+ * 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 migration 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";
+            std::cout << sim.get_result()[0].transpose() << "\n";
+            std::cout << sim.get_result().get_last_value().transpose() << "\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();
diff --git a/cpp/memilio/utils/random_number_generator.h b/cpp/memilio/utils/random_number_generator.h
@@ -666,6 +666,13 @@ using UniformIntDistribution = DistributionAdapter<std::uniform_int_distribution
 template <class Real>
 using UniformDistribution = DistributionAdapter<std::uniform_real_distribution<Real>>;
 
+/**
+ * adapted poisson_distribution.
+ * @see DistributionAdapter
+ */
+template <class Int>
+using PoissonDistribution = DistributionAdapter<std::poisson_distribution<Int>>;
+
 } // namespace mio
 
 #endif
diff --git a/cpp/memilio/utils/stl_util.h b/cpp/memilio/utils/stl_util.h
@@ -22,6 +22,8 @@
 
 #include "memilio/utils/metaprogramming.h"
 
+#include <array>
+#include <numeric>
 #include <vector>
 #include <algorithm>
 #include <utility>
@@ -289,6 +291,30 @@ bool contains(Iter b, Iter e, Pred p)
     return find_if(b, e, p) != e;
 }
 
+/**
+ * Get an std::array that contains all members of an enum class.
+ * The enum class must be a valid index, i.e. members must be sequential starting at 0 and there must
+ * be a member `Count` at the end, that will not be included in the array.
+ * Example:
+ * ```
+ * enum class E { A, B, Count };
+ * assert(enum_members<E>() == std::array<2, E>(E::A, E::B));
+ * ``` 
+ * @tparam T An enum class that is a valid index.
+ * @return Array of all members of the enum class not including T::Count.
+ */
+template<class T>
+constexpr std::array<T, size_t(T::Count)> enum_members()
+{
+    auto enum_members = std::array<T, size_t(T::Count)>{};
+    auto indices      = std::array<std::underlying_type_t<T>, size_t(T::Count)>{};
+    std::iota(indices.begin(), indices.end(), std::underlying_type_t<T>(0));
+    std::transform(indices.begin(), indices.end(), enum_members.begin(), [](auto i) {
+        return T(i);
+    });
+    return enum_members;
+}
+
 } // namespace mio
 
 #endif //STL_UTIL_H
diff --git a/cpp/tests/test_stl_util.cpp b/cpp/tests/test_stl_util.cpp
@@ -243,6 +243,17 @@ TEST(TestContains, empty)
     }));
 }
 
+TEST(EnumMembers, works)
+{
+    enum class E
+    {
+        A,
+        B,
+        Count
+    };
+    ASSERT_THAT(mio::enum_members<E>(), testing::ElementsAre(E::A, E::B));
+}
+
 TEST(TestContains, set_ostream_format)
 {
     std::ostringstream output;
