diff --git a/cpp/examples/graph_abm.cpp b/cpp/examples/graph_abm.cpp index 9eb1caf65a..aca8763d65 100644 --- a/cpp/examples/graph_abm.cpp +++ b/cpp/examples/graph_abm.cpp @@ -51,17 +51,29 @@ struct Logger : mio::LogAlways { { Type location_information{}; location_information.reserve(size_t(mio::abm::LocationType::Count)); - auto t = sim.get_time(); - for (auto&& loc : sim.get_model().get_locations()) { - std::map persons_per_infection_state; + auto t = sim.get_time(); + const auto& model = sim.get_model(); + const auto num_loc = model.get_locations().size(); + + // Count the Persons per Location and InfectionState in a single pass over all Persons, instead of calling + // get_subpopulation for every Location and InfectionState, which iterates over all Persons each time. + std::vector> persons_per_infection_state(num_loc); + for (auto& counts : persons_per_infection_state) { for (size_t i = 0; i < static_cast(mio::abm::InfectionState::Count); ++i) { - auto inf_state = mio::abm::InfectionState(i); - persons_per_infection_state.insert( - {inf_state, sim.get_model().get_subpopulation(loc.get_id(), t, inf_state)}); + counts[mio::abm::InfectionState(i)] = 0; + } + } + for (auto&& person : model.get_persons()) { + // Persons currently located in another model are counted there + if (person.get_location_model_id() == model.get_id()) { + ++persons_per_infection_state[person.get_location().get()][person.get_infection_state(t)]; } + } + + for (auto&& loc : model.get_locations()) { location_information.push_back(std::make_tuple(loc.get_model_id(), loc.get_type(), loc.get_id(), - sim.get_model().get_number_persons(loc.get_id()), - persons_per_infection_state)); + model.get_number_persons(loc.get_id()), + persons_per_infection_state[loc.get_id().get()])); } return location_information; } diff --git a/cpp/models/abm/model.cpp b/cpp/models/abm/model.cpp index a886edde98..0b1eac064a 100755 --- a/cpp/models/abm/model.cpp +++ b/cpp/models/abm/model.cpp @@ -380,18 +380,18 @@ LocationId Model::find_location(LocationType type, const PersonId person) const size_t Model::get_subpopulation_combined(TimePoint t, InfectionState s) const { - return std::accumulate(m_locations.begin(), m_locations.end(), (size_t)0, - [t, s, this](size_t running_sum, const Location& loc) { - return running_sum + get_subpopulation(loc.get_id(), t, s); - }); + // a single pass over all Persons, instead of calling get_subpopulation for every Location + return std::count_if(m_persons.begin(), m_persons.end(), [&](auto&& p) { + return p.get_location_model_id() == m_id && p.get_infection_state(t) == s; + }); } size_t Model::get_subpopulation_combined_per_location_type(TimePoint t, InfectionState s, LocationType type) const { - return std::accumulate( - m_locations.begin(), m_locations.end(), (size_t)0, [t, s, type, this](size_t running_sum, const Location& loc) { - return loc.get_type() == type ? running_sum + get_subpopulation(loc.get_id(), t, s) : running_sum; - }); + // a single pass over all Persons, instead of calling get_subpopulation for every Location + return std::count_if(m_persons.begin(), m_persons.end(), [&](auto&& p) { + return p.get_location_model_id() == m_id && p.get_location_type() == type && p.get_infection_state(t) == s; + }); } TripList& Model::get_trip_list() diff --git a/cpp/models/abm/model.h b/cpp/models/abm/model.h index ee9077932f..0feeb02879 100644 --- a/cpp/models/abm/model.h +++ b/cpp/models/abm/model.h @@ -394,6 +394,8 @@ class Model /** * @brief Get the number of Person%s of a particular #InfectionState for all Cell%s. + * Note that this iterates over all Person%s of the Model. To count the Person%s of many Location%s or + * #InfectionState%s, do not call this in a loop, but count them in a single pass over get_persons() instead. * @param[in] location A LocationId from the Model. * @param[in] t TimePoint of querry. * @param[in] state #InfectionState of interest. diff --git a/cpp/tests/test_abm_model.cpp b/cpp/tests/test_abm_model.cpp index fe7926ff06..6522f78f56 100644 --- a/cpp/tests/test_abm_model.cpp +++ b/cpp/tests/test_abm_model.cpp @@ -178,6 +178,14 @@ TEST_F(TestModel, getSubpopulationCombined) 2); // Verify the total count of persons with no symptoms across all locations. EXPECT_EQ(model.get_subpopulation_combined(t, mio::abm::InfectionState::InfectedNoSymptoms), 3); + + // A person currently located in another model (e.g. in a graph ABM) is not counted. + auto visitor = add_test_person(model, school1, age_group_15_to_34, mio::abm::InfectionState::InfectedNoSymptoms); + model.get_person(visitor).set_location(mio::abm::LocationType::School, school1, model.get_id() + 1); + EXPECT_EQ(model.get_subpopulation_combined_per_location_type(t, mio::abm::InfectionState::InfectedNoSymptoms, + mio::abm::LocationType::School), + 2); + EXPECT_EQ(model.get_subpopulation_combined(t, mio::abm::InfectionState::InfectedNoSymptoms), 3); } /**