From ac417d5c4e5e52795721e78491efc1a50e926da4 Mon Sep 17 00:00:00 2001
From: Vaclav Petras <wenzeslaus@gmail.com>
Date: Thu, 2 Nov 2023 08:54:18 -0400
Subject: [PATCH] Add mortality and survival rate actions

---
 include/pops/actions.hpp         | 15 ++-------------
 include/pops/host_pool.hpp       | 21 +++++++++++++++++++++
 include/pops/model.hpp           |  4 ++--
 include/pops/multi_host_pool.hpp | 23 +++++++++++++++++++++++
 4 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/include/pops/actions.hpp b/include/pops/actions.hpp
index 3b064847..33d0603c 100644
--- a/include/pops/actions.hpp
+++ b/include/pops/actions.hpp
@@ -225,19 +225,8 @@ class SurvivalRateAction
             int i = indices[0];
             int j = indices[1];
             if (survival_rate_(i, j) < 1) {
-                // remove percentage of infestation/infection in the infected class
-                auto infected = hosts.infected_at(i, j);
-                int removed_infected =
-                    infected - std::lround(infected * survival_rate_(i, j));
-                hosts.remove_infected_at(
-                    i, j, removed_infected, generator.survival_rate());
-                // remove the same percentage for total exposed and remove randomly from
-                // each cohort
-                auto exposed = hosts.exposed_at(i, j);
-                int total_removed_exposed =
-                    exposed - std::lround(exposed * survival_rate_(i, j));
-                hosts.remove_exposed_at(
-                    i, j, total_removed_exposed, generator.survival_rate());
+                hosts.remove_infection_by_ratio_at(
+                    i, j, survival_rate_(i, j), generator.survival_rate());
             }
         }
     }
diff --git a/include/pops/host_pool.hpp b/include/pops/host_pool.hpp
index 51232900..6acafa4a 100644
--- a/include/pops/host_pool.hpp
+++ b/include/pops/host_pool.hpp
@@ -562,6 +562,27 @@ class HostPool : public HostPoolInterface<RasterIndex>
         this->remove_infected_at(row, col, count, generator);
     }
 
+    /**
+     * @brief Remove percentage of infestation/infection
+     *
+     * remove the same percentage for total exposed and remove randomly from each cohort
+     *
+     * @param row Row index of the cell
+     * @param col Column index of the cell
+     * @param ratio Ratio of removed infection
+     * @param generator Random number generator to provide stochasticity for mortality
+     */
+    void remove_infection_by_ratio_at(
+        RasterIndex row, RasterIndex col, double ratio, Generator& generator)
+    {
+        auto infected = this->infected_at(row, col);
+        int removed_infected = infected - std::lround(infected * ratio);
+        this->remove_infected_at(row, col, removed_infected, generator);
+        auto exposed = this->exposed_at(row, col);
+        int total_removed_exposed = exposed - std::lround(exposed * ratio);
+        this->remove_exposed_at(row, col, total_removed_exposed, generator);
+    }
+
     /**
      * @brief Remove exposed hosts and make the hosts susceptible
      *
diff --git a/include/pops/model.hpp b/include/pops/model.hpp
index 440387f1..63263200 100644
--- a/include/pops/model.hpp
+++ b/include/pops/model.hpp
@@ -364,9 +364,9 @@ class Model
             // expectation is that mortality tracker is of length (1/mortality_rate
             // + mortality_time_lag).
             // TODO: died.zero(); should be done by the caller if needed, document!
-            Mortality<StandardHostPool, IntegerRaster, FloatRaster> mortality(
+            Mortality<StandardMultiHostPool, IntegerRaster, FloatRaster> mortality(
                 config_.mortality_rate, config_.mortality_time_lag);
-            mortality.action(host_pool);
+            mortality.action(multi_host_pool);
         }
         // compute spread rate
         if (config_.use_spreadrates && config_.spread_rate_schedule()[step]) {
diff --git a/include/pops/multi_host_pool.hpp b/include/pops/multi_host_pool.hpp
index 5e80718d..468dd8c7 100644
--- a/include/pops/multi_host_pool.hpp
+++ b/include/pops/multi_host_pool.hpp
@@ -64,6 +64,29 @@ class MultiHostPool
         }
     }
 
+    void remove_infection_by_ratio_at(
+        RasterIndex row, RasterIndex col, double ratio, Generator& generator)
+    {
+        for (auto& host_pool : host_pools_) {
+            host_pool->remove_infection_by_ratio_at(row, col, ratio, generator);
+        }
+    }
+
+    void apply_mortality_at(
+        RasterIndex row, RasterIndex col, double mortality_rate, int mortality_time_lag)
+    {
+        for (auto& host_pool : host_pools_) {
+            host_pool->apply_mortality_at(row, col, mortality_rate, mortality_time_lag);
+        }
+    }
+
+    void step_forward_mortality()
+    {
+        for (auto& host_pool : host_pools_) {
+            host_pool->step_forward_mortality();
+        }
+    }
+
     bool do_establishment_test(double value)
     {
         UNUSED(value);