Working on distribution functions

UofUEpiBio · Jun 11, 2024 · 2050a55 · 2050a55
1 parent c22f77b
commit 2050a55
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 6 deletions.
diff --git a/include/epiworld/entity-meat.hpp b/include/epiworld/entity-meat.hpp
@@ -17,23 +17,23 @@ inline EntityToAgentFun<TSeq> entity_to_unassigned_agents()
 
         // Figuring out how many to sample
         int n_to_sample;
-        if (e.prevalence_as_proportion)
+        if (e.get_prevalence_as_proportion())
         {
-            n_to_sample = static_cast<int>(std::floor(e.prevalence * n));
+            n_to_sample = static_cast<int>(std::floor(e.get_prevalence() * n));
             if (n_to_sample > static_cast<int>(n))
                 --n_to_sample;
 
         } else
         {
-            n_to_sample = static_cast<int>(e.prevalence);
+            n_to_sample = static_cast<int>(e.get_prevalence());
             if (n_to_sample > static_cast<int>(n))
                 throw std::range_error("There are only " + std::to_string(n) + 
                 " individuals in the population. Cannot add the entity to " +
                     std::to_string(n_to_sample));
         }
 
         int n_left = n;
-        for (size_t i = 0u; i < n_to_sample; ++i)
+        for (int i = 0; i < n_to_sample; ++i)
         {
             int loc = static_cast<epiworld_fast_uint>(
                 floor(m->runif() * n_left--)
@@ -358,7 +358,7 @@ inline void Entity<TSeq>::distribute()
 
         int n_left = n;
         std::iota(idx.begin(), idx.end(), 0);
-        while (n_to_assign > 0)
+        while ((n_to_assign > 0) && (n_left > 0))
         {
             int loc = static_cast<epiworld_fast_uint>(
                 floor(model->runif() * n_left--)
@@ -371,9 +371,12 @@ inline void Entity<TSeq>::distribute()
             auto & agent = model->get_agent(idx[loc]);
 
             if (!agent.has_entity(id))
+            {
                 agent.add_entity(
                     *this, this->model, this->state_init, this->queue_init
                     );
+                n_to_assign--;
+            }
 
             std::swap(idx[loc], idx[n_left]);
 

diff --git a/tests/05-mixing.cpp b/tests/05-mixing.cpp
@@ -134,6 +134,10 @@ EPIWORLD_TEST_CASE("SEIRMixing", "[SEIR-mixing]") {
     e1.set_dist_fun(nullptr);
     e2.set_dist_fun(nullptr);
     e3.set_dist_fun(nullptr);
+
+    e1.set_prevalence(2000, false);
+    e2.set_prevalence(2000, false);
+    e3.set_prevalence(2000, false);
 
     model.rm_entity(0);
     model.rm_entity(1);
@@ -146,6 +150,33 @@ EPIWORLD_TEST_CASE("SEIRMixing", "[SEIR-mixing]") {
     // Running and checking the results
     model.run(50, 123);
 
+    auto agents1 = model.get_entity(0).get_agents();
+    auto agents2 = model.get_entity(1).get_agents();
+    auto agents3 = model.get_entity(2).get_agents();
+
+    std::vector< int > counts(model.size(), 0);
+    for (const auto & a: agents1)
+        counts[a]++;
+
+    for (const auto & a: agents2)
+        counts[a]++;
+
+    for (const auto & a: agents3)
+        counts[a]++;
+
+    double n0 = 0, n1 = 0, n2 = 0, n3 = 0;
+    for (const auto & c: counts)
+    {
+        if (c == 0)
+            n0++;
+        else if (c == 1)
+            n1++;
+        else if (c == 2)
+            n2++;
+        else if (c == 3)
+            n3++;
+    }
+
 
     #ifndef CATCH_CONFIG_MAIN
     return 0;

diff --git a/tests/07-entitifuns.cpp b/tests/07-entitifuns.cpp
@@ -57,7 +57,7 @@ EPIWORLD_TEST_CASE("Entity member", "[Entity]") {
 
     #ifdef CATCH_CONFIG_MAIN
     REQUIRE(nrepeats == 0);
-    REQUIRE(ntrue == model.size());
+    REQUIRE(ntrue == static_cast<int>(model.size()));
     #endif
 
     // Distributing entities among agents in a small population: