diff --git a/docs/source/usage/how_to_run.rst b/docs/source/usage/how_to_run.rst index e2a9e31..3db0781 100644 --- a/docs/source/usage/how_to_run.rst +++ b/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``) diff --git a/examples/inputs.defaults b/examples/inputs.defaults index 2840ff2..7d506a6 100644 --- a/examples/inputs.defaults +++ b/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 diff --git a/src/AgentDefinitions.H b/src/AgentDefinitions.H index 6157b5a..e2d9acb 100644 --- a/src/AgentDefinitions.H +++ b/src/AgentDefinitions.H @@ -5,6 +5,7 @@ #ifndef _AGENT_DEF_H_ #define _AGENT_DEF_H_ +#include "AgentEnums.H" #include 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 AMREX_GPU_DEVICE AMREX_FORCE_INLINE void setAgentData (PTDType& ptd, int ip, int i, int j, const amrex::GpuArray& 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((i + 0.5_rt) * dx[0]); @@ -344,9 +235,20 @@ void setAgentData (PTDType& ptd, int ip, int i, int j, const amrex::GpuArrayinitial_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::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; diff --git a/src/AgentEnums.H b/src/AgentEnums.H new file mode 100644 index 0000000..f9de703 --- /dev/null +++ b/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 + +/*! \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 diff --git a/src/CensusData.H b/src/CensusData.H index 98bc4bf..b1ed317 100644 --- a/src/CensusData.H +++ b/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& dx, amrex::Long id, int cpu, int age_group, int family, int nborhood, int n_disease, amrex::Array4 const& nr_arr, amrex::Array4 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); diff --git a/src/CensusData.cpp b/src/CensusData.cpp index 6b97276..a3ae943 100644 --- a/src/CensusData.cpp +++ b/src/CensusData.cpp @@ -254,6 +254,14 @@ void CensusData::initAgents (AgentContainer& pc, /*!< Agents */ auto my_proc = ParallelDescriptor::MyProc(); int n_disease = pc.m_num_diseases; + // Get disease parameters for GPU - use AsyncArray for device access + Gpu::AsyncArray disease_parms_arr(n_disease); + auto* disease_parms_ptr = disease_parms_arr.data(); + for (int d = 0; d < n_disease; d++) { + disease_parms_ptr[d] = pc.getDiseaseParameters_d(d); + } + const DiseaseParm** disease_parms_d = disease_parms_ptr; + Long pid; #ifdef AMREX_USE_OMP #pragma omp critical(init_agents_nextid) @@ -312,7 +320,7 @@ void CensusData::initAgents (AgentContainer& pc, /*!< Agents */ age_group = AgeGroups::a18to29; /* single adult age 19-29 */ } setAgentDataAndAssignSchool(ptd, ip, i, j, k, dx, pid + ip, my_proc, age_group, family, nborhood, n_disease, - nr_arr, student_counts_arr, engine); + nr_arr, student_counts_arr, disease_parms_d, engine); } else if (family_size == 2) { if (il2 == 0) { /* 1% probability of one parent + one child */ @@ -327,14 +335,14 @@ void CensusData::initAgents (AgentContainer& pc, /*!< Agents */ age_group = AgeGroups::a18to29; /* one parent 19-29 */ } setAgentDataAndAssignSchool(ptd, ip, i, j, k, dx, pid + ip, my_proc, age_group, family, nborhood, - n_disease, nr_arr, student_counts_arr, engine); + n_disease, nr_arr, student_counts_arr, disease_parms_d, engine); if (((int)Random_int(100, engine)) < p_schoolage) { age_group = AgeGroups::a5to17; } else { age_group = AgeGroups::u5; } setAgentDataAndAssignSchool(ptd, ip + 1, i, j, k, dx, pid + ip + 1, my_proc, age_group, family, nborhood, - n_disease, nr_arr, student_counts_arr, engine); + n_disease, nr_arr, student_counts_arr, disease_parms_d, engine); } else { /* 2 adults, 28% over 65 (ASSUME both same age group) */ if (il2 < 28) { @@ -347,9 +355,9 @@ void CensusData::initAgents (AgentContainer& pc, /*!< Agents */ age_group = AgeGroups::a18to29; /* single adult age 19-29 */ } setAgentDataAndAssignSchool(ptd, ip, i, j, k, dx, pid + ip, my_proc, age_group, family, nborhood, - n_disease, nr_arr, student_counts_arr, engine); + n_disease, nr_arr, student_counts_arr, disease_parms_d, engine); setAgentDataAndAssignSchool(ptd, ip + 1, i, j, k, dx, pid + ip + 1, my_proc, age_group, family, nborhood, - n_disease, nr_arr, student_counts_arr, engine); + n_disease, nr_arr, student_counts_arr, disease_parms_d, engine); } } @@ -365,9 +373,9 @@ void CensusData::initAgents (AgentContainer& pc, /*!< Agents */ age_group = AgeGroups::a18to29; /* parents 19-29 */ } setAgentDataAndAssignSchool(ptd, ip, i, j, k, dx, pid + ip, my_proc, age_group, family, nborhood, n_disease, - nr_arr, student_counts_arr, engine); + nr_arr, student_counts_arr, disease_parms_d, engine); setAgentDataAndAssignSchool(ptd, ip + 1, i, j, k, dx, pid + ip + 1, my_proc, age_group, family, nborhood, - n_disease, nr_arr, student_counts_arr, engine); + n_disease, nr_arr, student_counts_arr, disease_parms_d, engine); /* Now pick the children's age groups */ for (int nc = 2; nc < family_size; ++nc) { @@ -377,7 +385,7 @@ void CensusData::initAgents (AgentContainer& pc, /*!< Agents */ age_group = AgeGroups::u5; } setAgentDataAndAssignSchool(ptd, ip + nc, i, j, k, dx, pid + ip + nc, my_proc, age_group, family, - nborhood, n_disease, nr_arr, student_counts_arr, engine); + nborhood, n_disease, nr_arr, student_counts_arr, disease_parms_d, engine); } } } diff --git a/src/DiseaseParm.H b/src/DiseaseParm.H index b7da51b..767cf30 100644 --- a/src/DiseaseParm.H +++ b/src/DiseaseParm.H @@ -6,10 +6,11 @@ #define DISEASE_PARM_H_ #include +#include #include #include -#include "AgentDefinitions.H" +#include "AgentEnums.H" using amrex::ParticleReal; using amrex::Real; @@ -38,6 +39,8 @@ struct DiseaseParm { int initial_case_type = CaseTypes::rnd; /*! Number of initial cases (in case of random initialization) */ int num_initial_cases = 0; + /*! Fraction of agents initially immune (0.0 to 1.0) */ + Real initial_immunity_fraction = Real(0.0); /*! Name of disease. This is a char array because this whole structure is memcpy'd to the device */ char disease_name[50]; /*! Initial cases filename (CaseData::initFromFile): diff --git a/src/DiseaseParm.cpp b/src/DiseaseParm.cpp index d8d149c..55941fc 100644 --- a/src/DiseaseParm.cpp +++ b/src/DiseaseParm.cpp @@ -29,6 +29,8 @@ void DiseaseParm::readInputs (const std::string& a_pp_str /*!< Parmparse string amrex::Abort("initial case type not recognized"); } + pp.query("initial_immunity_fraction", initial_immunity_fraction); + queryArray(pp, "xmit_comm", xmit_comm, AgeGroups::total); queryArray(pp, "xmit_hood", xmit_hood, AgeGroups::total); queryArray(pp, "xmit_hh_adult", xmit_hh_adult, AgeGroups::total); diff --git a/src/UrbanPopData.cpp b/src/UrbanPopData.cpp index 23d6f8a..162f89c 100644 --- a/src/UrbanPopData.cpp +++ b/src/UrbanPopData.cpp @@ -365,15 +365,22 @@ void UrbanPopData::initAgents (AgentContainer& pc, const ExaEpi::TestParams& par int i_RT = IntIdx::nattribs; int r_RT = RealIdx::nattribs; int n_disease = pc.m_num_diseases; + + // Get disease parameters for GPU - use AsyncArray for device access + Gpu::AsyncArray disease_parms_arr(n_disease); + auto* disease_parms_ptr = disease_parms_arr.data(); + for (int d = 0; d < n_disease; d++) { + disease_parms_ptr[d] = pc.getDiseaseParameters_d(d); + } + const DiseaseParm** disease_parms_d = disease_parms_ptr; + for (int d = 0; d < n_disease; d++) { soa.GetRealData(r_RT + r0(d) + RealIdxDisease::treatment_timer).assign(0.0_rt); - soa.GetRealData(r_RT + r0(d) + RealIdxDisease::disease_counter).assign(0.0_rt); soa.GetRealData(r_RT + r0(d) + RealIdxDisease::prob).assign(0.0_rt); soa.GetRealData(r_RT + r0(d) + RealIdxDisease::latent_period).assign(0.0_rt); soa.GetRealData(r_RT + r0(d) + RealIdxDisease::infectious_period).assign(0.0_rt); soa.GetRealData(r_RT + r0(d) + RealIdxDisease::incubation_period).assign(0.0_rt); soa.GetRealData(r_RT + r0(d) + RealIdxDisease::hospital_delay).assign(0.0_rt); - soa.GetIntData(i_RT + i0(d) + IntIdxDisease::status).assign(0); soa.GetIntData(i_RT + i0(d) + IntIdxDisease::symptomatic).assign(0); } auto np = soa.numParticles(); @@ -457,6 +464,36 @@ void UrbanPopData::initAgents (AgentContainer& pc, const ExaEpi::TestParams& par }); Gpu::synchronize(); + // Set initial disease status and immunity - use AsyncArray for device access + Gpu::AsyncArray status_ptrs_arr(n_disease); + Gpu::AsyncArray disease_counter_ptrs_arr(n_disease); + auto* status_ptrs = status_ptrs_arr.data(); + auto* disease_counter_ptrs = disease_counter_ptrs_arr.data(); + for (int d = 0; d < n_disease; d++) { + status_ptrs[d] = soa.GetIntData(i_RT + i0(d) + IntIdxDisease::status).data(); + disease_counter_ptrs[d] = soa.GetRealData(r_RT + r0(d) + RealIdxDisease::disease_counter).data(); + } + + ParallelForRNG(np, [=] AMREX_GPU_DEVICE (int i, RandomEngine const& engine) noexcept { + for (int d = 0; d < n_disease; d++) { + // Check if agent should be initially immune for this disease + if (disease_parms_d[d]->initial_immunity_fraction > 0.0_prt && + amrex::Random(engine) < disease_parms_d[d]->initial_immunity_fraction) { + status_ptrs[d][i] = Status::immune; + // Set immune counter to random point in immunity period + // Sample full immune duration, then pick random point within it + ParticleReal full_immune_duration = + static_cast(amrex::RandomGamma(disease_parms_d[d]->immune_length_alpha, + disease_parms_d[d]->immune_length_beta, engine)); + disease_counter_ptrs[d][i] = amrex::Random(engine) * full_immune_duration; + } else { + status_ptrs[d][i] = Status::never; + disease_counter_ptrs[d][i] = 0.0_prt; + } + } + }); + Gpu::synchronize(); + // now ensure that all members of the same family have the same home nborhood // and ensure all members of the same hh cluster have the same home neighborhood ParallelFor(np, [=] AMREX_GPU_DEVICE (int i) noexcept {