Implement start and stop for shelter-in-place (#36)

* implement start and stop for shelter-in-place

* remove unused

* document new parameters

docs/source/usage/how_to_run.rst

@@ -64,6 +64,12 @@ In addition to the ExaEpi inputs, there are also a number of runtime options tha
     aggregated data files will be named `cases000010`, etc.
 * ``agent.seed`` (`long integer`)
     Use this to specify the random seed to use for the run.
+* ``agent.shelter_start`` (`integer`)
+    Day on which to start shelter-in-place.
+* ``agent.shelter_length`` (`integer`)
+    Number of days shelter in-place-is in effect.
+* ``agent.symptomatic_withdraw`` (`integer`)
+    Whether or not to have symptomatic agents withdraw.
 * ``contact.pSC`` (`float`, default: 0.2)
     This is contact matrix scaling factor for schools.
 * ``contact.pCO`` (`float`, default: 1.45)

src/AgentContainer.H

@@ -193,6 +193,10 @@ public:
 
     void infectAgents ();
 
+    void shelterStart ();
+
+    void shelterStop ();
+
     void generateCellData (amrex::MultiFab& mf) const;
 
     std::array<amrex::Long, 5> getTotals ();

src/AgentContainer.cpp

@@ -1039,6 +1039,69 @@ void AgentContainer::interactAgents ()
     }
 }
 
+void AgentContainer::shelterStart ()
+{
+    BL_PROFILE("AgentContainer::shelterStart");
+
+    amrex::Print() << "Starting shelter in place order \n";
+
+    for (int lev = 0; lev <= finestLevel(); ++lev)
+    {
+        auto& plev  = GetParticles(lev);
+
+#ifdef AMREX_USE_OMP
+#pragma omp parallel if (Gpu::notInLaunchRegion())
+#endif
+        for(MFIter mfi = MakeMFIter(lev, TilingIfNotGPU()); mfi.isValid(); ++mfi)
+        {
+            int gid = mfi.index();
+            int tid = mfi.LocalTileIndex();
+            auto& ptile = plev[std::make_pair(gid, tid)];
+            auto& soa   = ptile.GetStructOfArrays();
+            const auto np = ptile.numParticles();
+            auto withdrawn_ptr = soa.GetIntData(IntIdx::withdrawn).data();
+
+            amrex::ParallelForRNG( np,
+            [=] AMREX_GPU_DEVICE (int i, amrex::RandomEngine const& engine) noexcept
+            {
+                if (amrex::Random(engine) < 0.95) {
+                    withdrawn_ptr[i] = 1;
+                }
+            });
+        }
+    }
+}
+
+void AgentContainer::shelterStop ()
+{
+    BL_PROFILE("AgentContainer::shelterStop");
+
+    amrex::Print() << "Stopping shelter in place order \n";
+
+    for (int lev = 0; lev <= finestLevel(); ++lev)
+    {
+        auto& plev  = GetParticles(lev);
+
+#ifdef AMREX_USE_OMP
+#pragma omp parallel if (Gpu::notInLaunchRegion())
+#endif
+        for(MFIter mfi = MakeMFIter(lev, TilingIfNotGPU()); mfi.isValid(); ++mfi)
+        {
+            int gid = mfi.index();
+            int tid = mfi.LocalTileIndex();
+            auto& ptile = plev[std::make_pair(gid, tid)];
+            auto& soa   = ptile.GetStructOfArrays();
+            const auto np = ptile.numParticles();
+            auto withdrawn_ptr = soa.GetIntData(IntIdx::withdrawn).data();
+
+            amrex::ParallelFor( np, [=] AMREX_GPU_DEVICE (int i) noexcept
+            {
+                withdrawn_ptr[i] = 0;
+            });
+        }
+    }
+}
+
 /*! \brief Infect agents based on their current status and the computed probability of infection.
     The infection probability is computed in AgentContainer::interactAgentsHomeWork() or
     AgentContainer::interactAgents()

src/Utils.H

@@ -65,6 +65,9 @@ struct TestParams
                                              (see: ExaEpi::IO::writeFIPSData) */
     std::string aggregated_diag_prefix; /*!< filename prefix for diagnostic data
                                              (see: ExaEpi::IO::writeFIPSData) */
+
+    int shelter_start = -1;
+    int shelter_length = 0;
 };
 
 /**

src/Utils.cpp

@@ -65,6 +65,9 @@ void ExaEpi::Utils::get_test_params (   TestParams& params,         /*!< Test pa
         pp.get("aggregated_diag_prefix", params.aggregated_diag_prefix);
     }
 
+    pp.query("shelter_start",  params.shelter_start);
+    pp.query("shelter_length", params.shelter_length);
+
     Long seed = 0;
     bool reset_seed = pp.query("seed", seed);
     if (reset_seed) {

src/main.cpp

@@ -191,6 +191,14 @@ void runAgent ()
                 ExaEpi::IO::writeFIPSData(pc, unit_mf, FIPS_mf, comm_mf, demo, params.aggregated_diag_prefix, i);
             }
 
+            if (params.shelter_start > 0 && params.shelter_start == i) {
+                pc.shelterStart();
+            }
+
+            if (params.shelter_start > 0 && params.shelter_start + params.shelter_length == i) {
+                pc.shelterStop();
+            }
+
             pc.updateStatus(disease_stats);
             pc.moveAgentsToWork();
             pc.interactAgentsHomeWork(mask_behavior, false);