Adding feature to seir conn

include/epiworld/models/seirconnected.hpp

@@ -4,6 +4,11 @@
 template<typename TSeq = EPI_DEFAULT_TSEQ>
 class ModelSEIRCONN : public epiworld::Model<TSeq> 
 {
+private:
+    std::vector< epiworld::Agent<TSeq> * > infected;
+    double effective_contact_rate;
+    void update_infected();
+
 public:
 
     static const int SUSCEPTIBLE = 0;
@@ -54,8 +59,35 @@ class ModelSEIRCONN : public epiworld::Model<TSeq>
         std::vector< int > queue_ = {}
     );
 
+    size_t get_n_infected() const { return infected.size(); }
+
 };
 
+template<typename TSeq>
+inline void ModelSEIRCONN<TSeq>::update_infected()
+{
+
+    infected.clear();
+    infected.reserve(this->size());
+
+    for (auto & p : this->get_agents())
+    {
+        if (p.get_state() == ModelSEIRCONN<TSeq>::INFECTED)
+        {
+            infected.push_back(&p);
+        }
+    }
+
+    Model<TSeq>::set_rand_binom(
+        this->get_n_infected(),
+        static_cast<double>(Model<TSeq>::par("Contact rate"))/
+            static_cast<double>(Model<TSeq>::size())
+    );
+
+    return;
+
+}
+
 template<typename TSeq>
 inline ModelSEIRCONN<TSeq> & ModelSEIRCONN<TSeq>::run(
     epiworld_fast_uint ndays,
@@ -74,13 +106,7 @@ inline void ModelSEIRCONN<TSeq>::reset()
 {
 
     Model<TSeq>::reset();
-
-    Model<TSeq>::set_rand_binom(
-        Model<TSeq>::size(),
-        static_cast<double>(
-            Model<TSeq>::par("Contact rate"))/
-            static_cast<double>(Model<TSeq>::size())
-        );
+    this->update_infected();
 
     return;
 
@@ -133,13 +159,16 @@ inline ModelSEIRCONN<TSeq>::ModelSEIRCONN(
             if (ndraw == 0)
                 return;
 
+            ModelSEIRCONN<TSeq> * model = dynamic_cast<ModelSEIRCONN<TSeq> *>(m);
+            size_t ninfected = model->get_n_infected();
+
             // Drawing from the set
             int nviruses_tmp = 0;
             for (int i = 0; i < ndraw; ++i)
             {
                 // Now selecting who is transmitting the disease
                 int which = static_cast<int>(
-                    std::floor(m->size() * m->runif())
+                    std::floor(ninfected * m->runif())
                 );
 
                 /* There is a bug in which runif() returns 1.0. It is rare, but
@@ -149,35 +178,32 @@ inline ModelSEIRCONN<TSeq>::ModelSEIRCONN(
                  * https://gcc.gnu.org/bugzilla/show_bug.cgi?id=63176
                  * 
                  */
-                if (which == static_cast<int>(m->size()))
+                if (which == static_cast<int>(ninfected))
                     --which;
 
+                epiworld::Agent<TSeq> & neighbor = *model->infected[which];
+
                 // Can't sample itself
-                if (which == static_cast<int>(p->get_id()))
+                if (neighbor.get_id() == p->get_id())
                     continue;
 
-                // If the neighbor is infected, then proceed
-                auto & neighbor = m->get_agents()[which];
-                if (neighbor.get_state() == ModelSEIRCONN<TSeq>::INFECTED)
-                {
+                // The neighbor is infected by construction
+                auto & v = neighbor.get_virus();
 
-                    auto & v = neighbor.get_virus();
-
-                    #ifdef EPI_DEBUG
-                    if (nviruses_tmp >= static_cast<int>(m->array_virus_tmp.size()))
-                        throw std::logic_error("Trying to add an extra element to a temporal array outside of the range.");
-                    #endif
-
-                    /* And it is a function of susceptibility_reduction as well */ 
-                    m->array_double_tmp[nviruses_tmp] =
-                        (1.0 - p->get_susceptibility_reduction(v, m)) * 
-                        v->get_prob_infecting(m) * 
-                        (1.0 - neighbor.get_transmission_reduction(v, m)) 
-                        ; 
-
-                    m->array_virus_tmp[nviruses_tmp++] = &(*v);
+                #ifdef EPI_DEBUG
+                if (nviruses_tmp >= static_cast<int>(m->array_virus_tmp.size()))
+                    throw std::logic_error("Trying to add an extra element to a temporal array outside of the range.");
+                #endif
+
+                /* And it is a function of susceptibility_reduction as well */ 
+                m->array_double_tmp[nviruses_tmp] =
+                    (1.0 - p->get_susceptibility_reduction(v, m)) * 
+                    v->get_prob_infecting(m) * 
+                    (1.0 - neighbor.get_transmission_reduction(v, m)) 
+                    ; 
+
+                m->array_virus_tmp[nviruses_tmp++] = &(*v);
 
-                }
             }
 
             // No virus to compute
@@ -279,6 +305,22 @@ inline ModelSEIRCONN<TSeq>::ModelSEIRCONN(
     model.add_state("Infected", update_infected);
     model.add_state("Recovered");
 
+    // Adding update function
+    epiworld::GlobalFun<TSeq> update = [](
+        epiworld::Model<TSeq> * m
+        ) -> void
+        {
+
+            ModelSEIRCONN<TSeq> * model = dynamic_cast<ModelSEIRCONN<TSeq> *>(m);
+
+            model->update_infected();
+
+            return;
+
+        };
+
+    model.add_globalevent(update, "Update infected individuals");
+
 
     // Preparing the virus -------------------------------------------
     epiworld::Virus<TSeq> virus(vname);

include/epiworld/models/sirconnected.hpp

@@ -189,7 +189,7 @@ inline ModelSIRCONN<TSeq>::ModelSIRCONN(
                 epiworld::Agent<TSeq> & neighbor = *model->infected[which];
 
                 // Can't sample itself
-                if (neighbor.get_id() == static_cast<int>(p->get_id()))
+                if (neighbor.get_id() == p->get_id())
                     continue;
 
                 // The neighbor is infected because it is on the list!