Created a function for setting an agent to infected, since that code …

…is used both during updates and when initializing agents to be infected.
AMReX-Codes · Nov 6, 2024 · 99c85ec · 99c85ec
1 parent f120169
commit 99c85ec
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 32 deletions.
diff --git a/examples/inputs.nm b/examples/inputs.nm
@@ -10,14 +10,14 @@ agent.urbanpop_filename = "../UrbanPop/NM_35/urbanpop_nm"
 agent.nsteps = 70
 #agent.seed = 11
 
-# disease.initial_case_type = "file"
-# disease.case_filename = "../CaseData/nm-july4.cases"
+disease.initial_case_type = "file"
+disease.case_filename = "../CaseData/nm-july4.cases"
 
-agent.number_of_diseases = 2
-agent.disease_names = "covid" "flu"
+# agent.number_of_diseases = 2
+# agent.disease_names = "covid" "flu"
 
-disease.initial_case_type = "random" "random"
-disease.num_initial_cases = 1000
+# disease.initial_case_type = "random" "random"
+# disease.num_initial_cases = 1000
 
 # setting these results in simulations with similar outcomes between UrbanPop and Census
 # disease.hospitalization_days = 6 6 6
@@ -28,8 +28,8 @@ disease.num_initial_cases = 1000
 # disease.icuCVF = 0 0 0 0 0 0.0
 # disease.ventCVF = 0.3 0.3 0.3 0.3 0.3 0.3
 
-disease_flu.num_initial_cases = 100
-disease_flu.p_trans = 0.3
-disease_flu.latent_length_mean = 3.0
-disease_flu.infectious_length_mean = 3.0
-disease_flu.incubation_length_mean = 3.0
+# disease_flu.num_initial_cases = 100
+# disease_flu.p_trans = 0.3
+# disease_flu.latent_length_mean = 3.0
+# disease_flu.infectious_length_mean = 3.0
+# disease_flu.incubation_length_mean = 3.0
diff --git a/src/AgentContainer.cpp b/src/AgentContainer.cpp
@@ -689,13 +689,13 @@ void AgentContainer::infectAgents ()
 
                 auto status_ptr = soa.GetIntData(i_RT+i0(d)+IntIdxDisease::status).data();
 
-                auto counter_ptr           = soa.GetRealData(r_RT+r0(d)+RealIdxDisease::disease_counter).data();
                 auto prob_ptr              = soa.GetRealData(r_RT+r0(d)+RealIdxDisease::prob).data();
+                auto counter_ptr           = soa.GetRealData(r_RT+r0(d)+RealIdxDisease::disease_counter).data();
                 auto latent_period_ptr     = soa.GetRealData(r_RT+r0(d)+RealIdxDisease::latent_period).data();
                 auto infectious_period_ptr = soa.GetRealData(r_RT+r0(d)+RealIdxDisease::infectious_period).data();
                 auto incubation_period_ptr = soa.GetRealData(r_RT+r0(d)+RealIdxDisease::incubation_period).data();
 
-                auto* lparm = m_d_parm[d];
+                const auto lparm = m_d_parm[d];
 
                 amrex::ParallelForRNG( np,
                 [=] AMREX_GPU_DEVICE (int i, amrex::RandomEngine const& engine) noexcept
@@ -704,17 +704,8 @@ void AgentContainer::infectAgents ()
                     if ( status_ptr[i] == Status::never ||
                          status_ptr[i] == Status::susceptible ) {
                         if (amrex::Random(engine) < prob_ptr[i]) {
-                            status_ptr[i] = Status::infected;
-                            counter_ptr[i] = 0.0_rt;
-                            latent_period_ptr[i] = amrex::RandomNormal(lparm->latent_length_mean, lparm->latent_length_std, engine);
-                            infectious_period_ptr[i] = amrex::RandomNormal(lparm->infectious_length_mean, lparm->infectious_length_std, engine);
-                            incubation_period_ptr[i] = amrex::RandomNormal(lparm->incubation_length_mean, lparm->incubation_length_std, engine);
-                            if (latent_period_ptr[i] < 0) { latent_period_ptr[i] = 0.0_rt; }
-                            if (infectious_period_ptr[i] < 0) { infectious_period_ptr[i] = 0.0_rt; }
-                            if (incubation_period_ptr[i] < 0) { incubation_period_ptr[i] = 0.0_rt; }
-                            if (incubation_period_ptr[i] > (infectious_period_ptr[i]+latent_period_ptr[i])) {
-                                incubation_period_ptr[i] = std::floor(infectious_period_ptr[i]+latent_period_ptr[i]);
-                            }
+                            setInfected(&(status_ptr[i]), &(counter_ptr[i]), &(latent_period_ptr[i]), &(infectious_period_ptr[i]),
+                                        &(incubation_period_ptr[i]), engine, lparm);
                             return;
                         }
                     }

diff --git a/src/AgentDefinitions.H b/src/AgentDefinitions.H
@@ -218,7 +218,7 @@ template <typename PTDType>
 AMREX_GPU_DEVICE AMREX_FORCE_INLINE
 bool notSusceptible ( const int      a_idx, /*!< Agent index */
                       const PTDType& a_ptd, /*!< Particle tile data */
-                     const int       a_d    /*!< Disease index */ )
+                      const int       a_d    /*!< Disease index */ )
 {
     return (   (a_ptd.m_runtime_idata[i0(a_d)+IntIdxDisease::status][a_idx] == Status::immune)
             || (a_ptd.m_runtime_idata[i0(a_d)+IntIdxDisease::status][a_idx] == Status::dead)
@@ -264,5 +264,4 @@ static int getSchoolType(const int grade) {
     return SchoolType::none;
 }
 
-
 #endif
diff --git a/src/DiseaseParm.H b/src/DiseaseParm.H
@@ -155,4 +155,30 @@ struct DiseaseParm
     }
 
 };
+
+
+/*! \brief Set this agent to infected status, and initialize disease periods. */
+AMREX_GPU_DEVICE AMREX_FORCE_INLINE
+void setInfected ( int* status,
+                   amrex::Real* counter,
+                   amrex::Real* latent_period,
+                   amrex::Real* infectious_period,
+                   amrex::Real* incubation_period,
+                   amrex::RandomEngine const& engine,
+                   const DiseaseParm* lparm)
+{
+    *status = Status::infected;
+    *counter = amrex::Real(0);
+    *latent_period = amrex::RandomNormal(lparm->latent_length_mean, lparm->latent_length_std, engine);
+    *infectious_period = amrex::RandomNormal(lparm->infectious_length_mean, lparm->infectious_length_std, engine);
+    *incubation_period = amrex::RandomNormal(lparm->incubation_length_mean, lparm->incubation_length_std, engine);
+    if (*latent_period < 0) { *latent_period = amrex::Real(0);}
+        if (*infectious_period < 0) { *infectious_period = amrex::Real(0);}
+    if (*incubation_period < 0) { *incubation_period = amrex::Real(0);}
+    if (*incubation_period > (*infectious_period + *latent_period)) {
+        *incubation_period = std::floor(*infectious_period + *latent_period);
+    }
+}
+
+
 #endif
diff --git a/src/InitializeInfections.cpp b/src/InitializeInfections.cpp
@@ -84,7 +84,7 @@ static int infect_random_community (AgentContainer& pc, /*!< Agent container (pa
 
         auto comm_arr = comm_mf[mfi].array();
 
-        const auto* lparm = pc.getDiseaseParameters_d(d_idx);
+        const auto lparm = pc.getDiseaseParameters_d(d_idx);
 
         Gpu::DeviceScalar<int> num_infected_d(num_infected);
         int* num_infected_p = num_infected_d.dataPtr();
@@ -112,11 +112,8 @@ static int infect_random_community (AgentContainer& pc, /*!< Agent container (pa
                         ip += ninfect;
                     }
                 } else {
-                    status_ptr[pindex] = Status::infected;
-                    counter_ptr[pindex] = 0;
-                    latent_period_ptr[pindex] = RandomNormal(lparm->latent_length_mean, lparm->latent_length_std, engine);
-                    infectious_period_ptr[pindex] = RandomNormal(lparm->infectious_length_mean, lparm->infectious_length_std, engine);
-                    incubation_period_ptr[pindex] = RandomNormal(lparm->incubation_length_mean, lparm->incubation_length_std, engine);
+                    setInfected(&(status_ptr[pindex]), &(counter_ptr[pindex]), &(latent_period_ptr[pindex]),
+                                &(infectious_period_ptr[pindex]), &(incubation_period_ptr[pindex]), engine, lparm);
                     ++ni;
                 }
             }