Add initial immunity fraction parameter

docs/source/usage/how_to_run.rst

@@ -128,6 +128,12 @@ The following inputs specify the disease parameters:
 * ``disease.num_initial_cases`` (`int`, default ``0``)
     The number of initial cases to seed for a single disease. Must be provided if
     ``initial_case_type`` is ``"random"``. It can be set to 0 for no cases.
+* ``disease.initial_immunity_fraction`` (`float`, default ``0.0``)
+    The fraction of agents (0.0 to 1.0) that are initially immune for each disease. This sets agents as immune at initialization.
+    Each initially immune agent is placed at a random point in their immunity period. The immunity period duration is sampled from a
+    Gamma distribution with the ``disease.immune_length_alpha`` and ``disease.immune_length_beta`` parameters, and the agent's
+    remaining immunity is set to a uniform random value between 0 and that duration. This feature works for both ``census`` and
+    ``urbanpop`` initialization types. For multiple diseases, use disease-specific parameters: ``disease_[disease name].initial_immunity_fraction``.
 * ``disease.p_trans`` (`float`, default ``0.2``)
     Probability of transmission given contact. There must be one entry for each disease strain.
 * ``disease.p_asymp`` (`float`, default ``0.4``)

examples/inputs.defaults

@@ -79,6 +79,15 @@ disease.initial_case_type = random
 disease.num_initial_cases = 0
 # no default, must be set for each disease when disease_[disease name].num_initial_cases = random for multiple diseases
 # disease.num_initial_cases_[disease name] = 0
+# The fraction of agents (0.0 to 1.0) that are initially immune for each disease.
+# This sets agents as immune at initialization. Each initially immune agent is placed at a random
+# point in their immunity period. The immunity period duration is sampled from a Gamma distribution
+# using the immune_length_alpha and immune_length_beta parameters, and the agent's remaining immunity
+# is set to a uniform random value between 0 and that duration.
+# This feature works for both Census and UrbanPop initialization types.
+disease.initial_immunity_fraction = 0.0
+# For multiple diseases, use disease-specific parameters:
+# disease_[disease name].initial_immunity_fraction = 0.0
 
 # Probability of transmission given contact.
 disease.p_trans = 0.2

src/AgentDefinitions.H

@@ -5,6 +5,7 @@
 #ifndef _AGENT_DEF_H_
 #define _AGENT_DEF_H_
 
+#include "AgentEnums.H"
 #include <AMReX_Particles.H>
 
 namespace ExaEpi {
@@ -19,72 +20,6 @@ struct RealIdx {
     };
 };
 
-/*! \brief Disease-specific Real-type Runtime-SoA attributes of agent */
-struct RealIdxDisease {
-    enum {
-        /* Disease counter starts after infection. */
-        treatment_timer = 0, /*!< Timer since hospital admission */
-        disease_counter,     /*!< Counter since start of infection */
-        prob,                /*!< Probability of infection */
-        latent_period,       /*!< Time until infectious, which could be before symptoms appear */
-        infectious_period,   /*!< Length of time infectious */
-        incubation_period,   /*!< Time until symptoms appear */
-        hospital_delay,      /*!< Delay after symptom onset until hospital treatment is sought */
-        hospital_random,     /*!< Random number that decides hospitaliation */
-        nattribs             /*!< number of real-type attribute*/
-    };
-};
-
-/*! \brief Integer-type SoA attributes of agent */
-struct IntIdx {
-    enum {
-        age_group = 0, /*!< Age group (under 5, 5-17, 18-29, 30-64, 65+) */
-        family,        /*!< Family ID */
-        home_i,        /*!< home location index */
-        home_j /*!< home location index */,
-        work_i /*!< work location index */,
-        work_j /*!< work location index */,
-        hosp_i /*!< hosp location index */,
-        hosp_j /*!< hosp location index */,
-        trav_i /*!< air travel location index */,
-        trav_j /*!< air travel location index */,
-        nborhood,      /*!< home neighborhood ID */
-        hh_cluster,    /*!< household cluster ID */
-        school_grade,  /*!< school grade, including universities */
-        school_id,     /*!< ID for a given school */
-        school_closed, /*!< 0 for open, 1 for closed */
-        naics,         /*!< industry NAICS code for business employed at */
-        workgroup,     /*!< workgroup ID */
-        work_nborhood, /*!< work neighborhood ID */
-        withdrawn,     /*!< quarantine status */
-        random_travel, /*!< on long distance travel? */
-        air_travel,    /*!< on long distance travel by Air? */
-        nattribs       /*!< number of integer-type attribute */
-    };
-};
-
-/*! \brief Disease-specific Integer-type Runtime-SoA attributes of agent */
-struct IntIdxDisease {
-    enum {
-        status = 0,  /*!< Disease status (#Status) */
-        symptomatic, /*!< currently symptomatic? 0: no, but will be, 1: yes, 2: no, and will remain so until recovered */
-        nattribs     /*!< number of integer-type attribute */
-    };
-};
-
-/*! \brief School Type  */
-struct SchoolType {
-    enum {
-        none, /*!< Not in school */
-        college,
-        high,
-        middle,
-        elem,
-        daycare,
-        total
-    };
-};
-
 /*! \brief School types used only in initializing census data approach */
 struct SchoolCensusIDType {
     enum {
@@ -97,54 +32,6 @@ struct SchoolCensusIDType {
         total
     };
 };
-
-/*! \brief Age Group */
-struct AgeGroups {
-    enum {
-        u5 = 0,  /*!< Under 5 */
-        a5to17,  /*!< 5-17 */
-        a18to29, /*!< 18-29 */
-        a30to49, /*!< 30-49 */
-        a50to64,
-        o65,     /*!< over 65 */
-        total    /*!< number of age groups */
-    };
-};
-
-/*! \brief Disease status */
-struct Status {
-    enum {
-        never = 0,   /*!< never infected */
-        infected,    /*!< infected */
-        immune,      /*!< no longer infected, immune. lasts 6 months. */
-        susceptible, /*!< no longer infected, no longer immnune */
-        dead         /*!< passed away */
-    };
-};
-
-/*! \brief Disease statistics */
-struct DiseaseStats {
-    enum {
-        hospitalization = 0, /*!< number of current hospitalizations */
-        ICU,                 /*!< number of current ICU cases */
-        ventilator,          /*!< number of current ventilator cases */
-        death,               /*!< number of deaths */
-        new_cases            /*!< number of new infections each day */
-    };
-};
-
-/*! \brief Compute index offsets for runtime int-type disease attributes */
-AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
-int i0 (const int a_d /*!< Disease index */) {
-    return a_d * IntIdxDisease::nattribs;
-}
-
-/*! \brief Compute index offsets for runtime real-type disease attributes */
-AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
-int r0 (const int a_d /*!< Disease index */) {
-    return a_d * RealIdxDisease::nattribs;
-}
-
 /*! \brief Disease symptom status */
 struct SymptomStatus {
     enum {
@@ -332,10 +219,14 @@ bool isNewlyHospitalized (const int a_idx,      /*!< Agent index */
     return disease_counter == incb_period + hospital_delay;
 }
 
+// Include DiseaseParm after basic definitions to avoid circular dependency
+#include "DiseaseParm.H"
+
 template <typename PTDType>
 AMREX_GPU_DEVICE AMREX_FORCE_INLINE
 void setAgentData (PTDType& ptd, int ip, int i, int j, const amrex::GpuArray<amrex::Real, 2>& dx, amrex::Long id, int cpu,
-                   int age_group, int family, int nborhood, int n_disease) {
+                   int age_group, int family, int nborhood, int n_disease, const DiseaseParm** disease_parms,
+                   const amrex::RandomEngine& engine) {
     using namespace amrex::literals;
 
     ptd[ip].pos(0) = static_cast<amrex::ParticleReal>((i + 0.5_rt) * dx[0]);
@@ -344,9 +235,20 @@ void setAgentData (PTDType& ptd, int ip, int i, int j, const amrex::GpuArray<amr
     ptd[ip].cpu() = cpu;
 
     for (int d = 0; d < n_disease; d++) {
-        ptd.m_runtime_idata[i0(d) + IntIdxDisease::status][ip] = 0;
+        // Check if agent should be initially immune for this disease
+        if (disease_parms[d]->initial_immunity_fraction > 0.0_prt &&
+            amrex::Random(engine) < disease_parms[d]->initial_immunity_fraction) {
+            ptd.m_runtime_idata[i0(d) + IntIdxDisease::status][ip] = Status::immune;
+            // Set immune counter to random point in immunity period
+            // Sample full immune duration, then pick random point within it
+            amrex::ParticleReal full_immune_duration = static_cast<amrex::ParticleReal>(
+                    amrex::RandomGamma(disease_parms[d]->immune_length_alpha, disease_parms[d]->immune_length_beta, engine));
+            ptd.m_runtime_rdata[r0(d) + RealIdxDisease::disease_counter][ip] = amrex::Random(engine) * full_immune_duration;
+        } else {
+            ptd.m_runtime_idata[i0(d) + IntIdxDisease::status][ip] = Status::never;
+            ptd.m_runtime_rdata[r0(d) + RealIdxDisease::disease_counter][ip] = 0.0_prt;
+        }
         ptd.m_runtime_rdata[r0(d) + RealIdxDisease::treatment_timer][ip] = 0.0_prt;
-        ptd.m_runtime_rdata[r0(d) + RealIdxDisease::disease_counter][ip] = 0.0_prt;
         ptd.m_runtime_rdata[r0(d) + RealIdxDisease::prob][ip] = 0.0_prt;
         ptd.m_runtime_rdata[r0(d) + RealIdxDisease::latent_period][ip] = 0.0_prt;
         ptd.m_runtime_rdata[r0(d) + RealIdxDisease::infectious_period][ip] = 0.0_prt;

src/AgentEnums.H

@@ -0,0 +1,123 @@
+/*! @file AgentEnums.H
+    \brief Contains basic enum definitions used across agent and disease code
+*/
+
+#ifndef _AGENT_ENUMS_H_
+#define _AGENT_ENUMS_H_
+
+#include <AMReX_GpuQualifiers.H>
+
+/*! \brief School Type  */
+struct SchoolType {
+    enum {
+        none, /*!< Not in school */
+        college,
+        high,
+        middle,
+        elem,
+        daycare,
+        total
+    };
+};
+
+/*! \brief Age Group */
+struct AgeGroups {
+    enum {
+        u5 = 0,  /*!< Under 5 */
+        a5to17,  /*!< 5-17 */
+        a18to29, /*!< 18-29 */
+        a30to49, /*!< 30-49 */
+        a50to64, /*!< 50-64 */
+        o65,     /*!< Over 65 */
+        total
+    };
+};
+
+/*! \brief Disease-specific Real-type Runtime-SoA attributes of agent */
+struct RealIdxDisease {
+    enum {
+        /* Disease counter starts after infection. */
+        treatment_timer = 0, /*!< Timer since hospital admission */
+        disease_counter,     /*!< Counter since start of infection */
+        prob,                /*!< Probability of infection */
+        latent_period,       /*!< Time until infectious, which could be before symptoms appear */
+        infectious_period,   /*!< Length of time infectious */
+        incubation_period,   /*!< Time until symptoms appear */
+        hospital_delay,      /*!< Delay after symptom onset until hospital treatment is sought */
+        hospital_random,     /*!< Random number that decides hospitaliation */
+        nattribs             /*!< number of real-type attribute*/
+    };
+};
+
+/*! \brief Integer-type SoA attributes of agent */
+struct IntIdx {
+    enum {
+        age_group = 0, /*!< Age group (under 5, 5-17, 18-29, 30-64, 65+) */
+        family,        /*!< Family ID */
+        home_i,        /*!< home location index */
+        home_j /*!< home location index */,
+        work_i /*!< work location index */,
+        work_j /*!< work location index */,
+        hosp_i /*!< hosp location index */,
+        hosp_j /*!< hosp location index */,
+        trav_i /*!< air travel location index */,
+        trav_j /*!< air travel location index */,
+        nborhood,      /*!< home neighborhood ID */
+        hh_cluster,    /*!< household cluster ID */
+        school_grade,  /*!< school grade, including universities */
+        school_id,     /*!< ID for a given school */
+        school_closed, /*!< 0 for open, 1 for closed */
+        naics,         /*!< industry NAICS code for business employed at */
+        workgroup,     /*!< workgroup ID */
+        work_nborhood, /*!< work neighborhood ID */
+        withdrawn,     /*!< quarantine status */
+        random_travel, /*!< on long distance travel? */
+        air_travel,    /*!< on long distance travel by Air? */
+        nattribs       /*!< number of integer-type attribute */
+    };
+};
+
+/*! \brief Disease-specific Integer-type Runtime-SoA attributes of agent */
+struct IntIdxDisease {
+    enum {
+        status = 0,  /*!< Disease status (#Status) */
+        symptomatic, /*!< currently symptomatic? 0: no, but will be, 1: yes, 2: no, and will remain so until recovered */
+        nattribs     /*!< number of integer-type attribute */
+    };
+};
+
+/*! \brief Disease status */
+struct Status {
+    enum {
+        never = 0,   /*!< Never infected */
+        infected,    /*!< Currently infected */
+        immune,      /*!< Immune (recovered or vaccinated) */
+        susceptible, /*!< Susceptible (waned immunity) */
+        dead         /*!< Dead */
+    };
+};
+
+/*! \brief Disease statistics */
+struct DiseaseStats {
+    enum {
+        hospitalization = 0, /*!< number of current hospitalizations */
+        ICU,                 /*!< number of current ICU cases */
+        ventilator,          /*!< number of current ventilator cases */
+        death,               /*!< number of deaths */
+        new_cases            /*!< number of new infections each day */
+    };
+};
+
+/*! \brief Compute index offsets for runtime integer-type disease attributes */
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+int i0 (const int a_d /*!< Disease index */) {
+    return a_d * IntIdxDisease::nattribs;
+}
+
+/*! \brief Compute index offsets for runtime real-type disease attributes */
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+int r0 (const int a_d /*!< Disease index */) {
+    return a_d * RealIdxDisease::nattribs;
+}
+
+#endif

src/CensusData.H

@@ -53,8 +53,8 @@ AMREX_GPU_DEVICE AMREX_FORCE_INLINE
 void setAgentDataAndAssignSchool (PTDType& ptd, int ip, int i, int j, int k, const amrex::GpuArray<amrex::Real, 2>& dx,
                                   amrex::Long id, int cpu, int age_group, int family, int nborhood, int n_disease,
                                   amrex::Array4<int> const& nr_arr, amrex::Array4<int> const& student_counts_arr,
-                                  const amrex::RandomEngine& engine) {
-    setAgentData(ptd, ip, i, j, dx, id, cpu, age_group, family, nborhood, n_disease);
+                                  const DiseaseParm** disease_parms, const amrex::RandomEngine& engine) {
+    setAgentData(ptd, ip, i, j, dx, id, cpu, age_group, family, nborhood, n_disease, disease_parms, engine);
     Gpu::Atomic::AddNoRet(&nr_arr(i, j, k, age_group), 1);
 
     assignSchool(&ptd.m_idata[IntIdx::school_grade][ip], &ptd.m_idata[IntIdx::school_id][ip], age_group, nborhood, engine);