WIP: Partial transmission file implementation

src/CMakeLists.txt

@@ -25,9 +25,10 @@ SET(LIBRARY_SRC
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/network.h
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output.h
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/OutputFile.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/PerformanceFile.h
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/PerIndividualFile.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/PerformanceFile.h
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/TimeSeriesFile.h
+    ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/TransmissionFile.h
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/person.h
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/population.h
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/random_interactions.h
@@ -43,10 +44,10 @@ SET(LIBRARY_SRC
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/input.cu
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/network.cu
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output.cu
-    ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/PerformanceFile.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/TimeSeriesFile.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/PerIndividualFile.cu
+    ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/PerformanceFile.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/TimeSeriesFile.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/output/TransmissionFile.cu
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/person.cu
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/population.cu
     ${CMAKE_CURRENT_SOURCE_DIR}/exateppabm/random_interactions.cu

src/exateppabm/cli.cpp

@@ -15,6 +15,7 @@ void setup(CLI::App& app, std::shared_ptr<exateppabm::cli::params> params) {
     app.add_option("-n, --param-number", params->inputParamLine, "The line from the parameters file to use. 1 indexed (assuming there is a header)");
     app.add_option("-o, --output-dir", params->outputDir, "Path to output directory");
     app.add_flag("--individual-file", params->individualFile, "Enable the creation of the per individual file");
+    app.add_flag("--transmission-file", params->transmissionFile, "Enable the transmission file. This has significant performance implications");
 }
 
 void print(const exateppabm::cli::params params) {
@@ -25,6 +26,7 @@ void print(const exateppabm::cli::params params) {
     fmt::print("  inputParamLine = {}\n", params.inputParamLine);
     fmt::print("  outputDir = {}\n", params.outputDir);
     fmt::print("  individualFile = {}\n", params.individualFile);
+    fmt::print("  transmissionFile = {}\n", params.transmissionFile);
     fmt::print("}}\n");
 }
 

src/exateppabm/cli.h

@@ -41,6 +41,10 @@ struct params {
      * bool indicating if the individual file should be written
      */
     bool individualFile = false;
+    /**
+     * bool indicating if the transmission file should be written
+     */
+    bool transmissionFile = false;
 };
 
 /**

src/exateppabm/disease/SEIR.cu

@@ -17,9 +17,6 @@ FLAMEGPU_AGENT_FUNCTION(progressDisease, flamegpu::MessageNone, flamegpu::Messag
     // Get the current timestep / day
     std::uint32_t today = FLAMEGPU->getStepCounter();
 
-    // Get a handle to the total_infected_per_demographic macro env property
-    auto totalInfectedPerDemographic = FLAMEGPU->environment.getMacroProperty<std::uint32_t, demographics::AGE_COUNT>("total_infected_per_demographic");
-
     // Get the current agents infection status
     auto infectionState = FLAMEGPU->getVariable<disease::SEIR::InfectionStateUnderlyingType>(person::v::INFECTION_STATE);
     // Get the time they last changed state
@@ -28,52 +25,23 @@ FLAMEGPU_AGENT_FUNCTION(progressDisease, flamegpu::MessageNone, flamegpu::Messag
     float stateDuration = FLAMEGPU->getVariable<float>(person::v::INFECTION_STATE_DURATION);
     // Ready to change state if today is past the next scheduled state change
     bool readyToChange = today >= dayOfLastStateChange + std::ceil(stateDuration);
-    // Get the agent's demographic
-    auto demographic_idx = FLAMEGPU->getVariable<demographics::AgeUnderlyingType>(person::v::AGE_DEMOGRAPHIC);
-
     // For each different initial state, change if required and compute the next state's duration.
-    if (infectionState == disease::SEIR::InfectionState::Susceptible) {
-        // no op
-    } else if (infectionState == disease::SEIR::InfectionState::Exposed) {
+    if (infectionState == disease::SEIR::InfectionState::Exposed) {
         // Exposed to Infected, if enough time has passed
         if (readyToChange) {
-            // Update the state
-            infectionState = disease::SEIR::InfectionState::Infected;
-            // Update the day
-            dayOfLastStateChange = today;
-            // Compute how long the next state will last
-            float mean = FLAMEGPU->environment.getProperty<float>("mean_time_to_recovered");
-            float sd = FLAMEGPU->environment.getProperty<float>("sd_time_to_recovered");
-            stateDuration = (FLAMEGPU->random.normal<float>() * sd) + mean;
-            // Atomically update the number of infected individuals  for the current individuals demographics when they transition into the infection state
-            totalInfectedPerDemographic[demographic_idx]++;
+            exposedToInfected(FLAMEGPU, infectionState);
         }
     } else if (infectionState == disease::SEIR::InfectionState::Infected) {
         // Infected to Recovered if enough time has passed
         if (readyToChange) {
-            // Update the state
-            infectionState = disease::SEIR::InfectionState::Recovered;
-            // Update the day
-            dayOfLastStateChange = today;
-            // Compute how long the next state will last
-            float mean = FLAMEGPU->environment.getProperty<float>("mean_time_to_susceptible");
-            float sd = FLAMEGPU->environment.getProperty<float>("sd_time_to_susceptible");
-            stateDuration = (FLAMEGPU->random.normal<float>() * sd) + mean;
+            infectedToRecovered(FLAMEGPU, infectionState);
         }
     } else if (infectionState == disease::SEIR::InfectionState::Recovered) {
         // Recovered to Susceptible, if enough time has passed.
         if (readyToChange) {
-            infectionState = disease::SEIR::Susceptible;
-            dayOfLastStateChange = today;
-            stateDuration = 0;  // susceptible doesn't have a fixed duration
+            recoveredToSusceptible(FLAMEGPU, infectionState);
         }
     }
-
-    // Update global agent variables from local (in register) values.
-    FLAMEGPU->setVariable<disease::SEIR::InfectionStateUnderlyingType>(person::v::INFECTION_STATE, infectionState);
-    FLAMEGPU->setVariable<std::uint32_t>(person::v::INFECTION_STATE_CHANGE_DAY, dayOfLastStateChange);
-    FLAMEGPU->setVariable<float>(person::v::INFECTION_STATE_DURATION, stateDuration);
-
     return flamegpu::ALIVE;
 }
 
@@ -86,7 +54,7 @@ void define(flamegpu::ModelDescription& model, const exateppabm::input::config&
     env.newMacroProperty<std::uint32_t, demographics::AGE_COUNT>("total_infected_per_demographic");
 
     // Add a number of model parameters to the environment, initialised with the value from the configuration file
-    // @todo - not all of these feel right here / add cosntexpr strings somewhere.
+    // @todo - not all of these feel right here / add constexpr strings somewhere.
     env.newProperty<float>("mean_time_to_infected", params.mean_time_to_infected);
     env.newProperty<float>("sd_time_to_infected", params.sd_time_to_infected);
     env.newProperty<float>("mean_time_to_recovered", params.mean_time_to_recovered);

src/exateppabm/disease/SEIR.h

@@ -4,13 +4,15 @@
 #include "flamegpu/flamegpu.h"
 #include "exateppabm/input.h"
 #include "exateppabm/person.h"
+#include "exateppabm/demographics.h"
+
 
 namespace exateppabm {
 namespace disease {
 
 /**
  * Methods and variables related to SEIR modelling, with no birth or death included
- * 
+ *
  * This is intended to demonstrate the type of model which can be implemented and how, with more complex disease modelling following a similar pattern.
  */
 namespace SEIR {
@@ -23,9 +25,9 @@ namespace SEIR {
 typedef std::uint32_t InfectionStateUnderlyingType;
 
 /**
- * Enum representing the different states of the model. 
+ * Enum representing the different states of the model.
  * Unfortunately unable to use a enum class for strong typing without excessive casting being required when passing to/from templated FLAMEGPU 2 methods.
- * 
+ *
  * @note - the static_casting required might be worthwhile if multiple infection models are implemented, as non-class enums are in the global scope without any prefixing.
  * @note - Possibly just define __host__ __device__ methods to/from_uint8() rather than having the cast? Not sure <type_traits> is available on device for a host device to_underlying pre c++23
  * @note - not just named "state" to reduce confusion with FLAME GPU 2 agent states.
@@ -40,7 +42,7 @@ enum InfectionState : InfectionStateUnderlyingType {
 
 /**
  * Add required parts to a FLAME GPU 2 model for this implementation.
- * 
+ *
  * @note - this is likely to be refactored several times
  * @param model the FLAME GPU Model description to mutate
  * @param parameters the model parameters
@@ -56,9 +58,9 @@ void appendLayers(flamegpu::ModelDescription& model);
 
 /**
  * Device utility function to get an individuals current infection status from global agent memory
- * 
+ *
  * Templated for due to the templated DeviceAPI object
- *  * 
+ *  *
  * @param FLAMEGPU flame gpu device API object
  * @return individuals current infection state
  */
@@ -70,11 +72,11 @@ FLAMEGPU_DEVICE_FUNCTION disease::SEIR::InfectionStateUnderlyingType getCurrentI
 
 /**
  * Device utility function for when an individual is exposed, moving from susceptible to exposed
- * 
+ *
  * Templated for due to the templated DeviceAPI object
- * 
+ *
  * @param FLAMEGPU flamegpu device api object
- * @param current infection status for the individual to be mutated in-place
+ * @param infectionStatus current infection status for the individual to be mutated in-place
  */
 template<typename MsgIn, typename MsgOut>
 FLAMEGPU_DEVICE_FUNCTION void susceptibleToExposed(flamegpu::DeviceAPI<MsgIn, MsgOut>* FLAMEGPU, disease::SEIR::InfectionStateUnderlyingType& infectionStatus) {
@@ -92,6 +94,119 @@ FLAMEGPU_DEVICE_FUNCTION void susceptibleToExposed(flamegpu::DeviceAPI<MsgIn, Ms
 
     // Increment the infection counter for this individual
     person::incrementInfectionCounter(FLAMEGPU);
+
+    // Update the agent variable tracking when the "current" infection exposure occurred.
+    FLAMEGPU->template setVariable<std::uint32_t>(person::v::TIME_EXPOSED, FLAMEGPU->getStepCounter());
+}
+
+
+/**
+ * Function for updating agent data when moved from the Exposed to Infected state
+ *
+ * Templated for due to the templated DeviceAPI object
+ *
+ * @param FLAMEGPU flamegpu device api object
+ * @param infectionStatus current infection status for the individual to be mutated in-place
+ */
+template<typename MsgIn, typename MsgOut>
+FLAMEGPU_DEVICE_FUNCTION void exposedToInfected(flamegpu::DeviceAPI<MsgIn, MsgOut>* FLAMEGPU, disease::SEIR::InfectionStateUnderlyingType& infectionStatus) {
+    std::uint32_t today = FLAMEGPU->getStepCounter();
+    // Get a handle to the total_infected_per_demographic macro env property
+    auto totalInfectedPerDemographic = FLAMEGPU->environment.template getMacroProperty<std::uint32_t, demographics::AGE_COUNT>("total_infected_per_demographic");
+
+    // Get the agent's demographic
+    auto demographic_idx = FLAMEGPU->template getVariable<demographics::AgeUnderlyingType>(person::v::AGE_DEMOGRAPHIC);
+
+    // Update the state
+    infectionStatus = disease::SEIR::InfectionState::Infected;
+    FLAMEGPU->template setVariable<disease::SEIR::InfectionStateUnderlyingType>(person::v::INFECTION_STATE, infectionStatus);
+
+    // Update the day
+    FLAMEGPU->template setVariable<std::uint32_t>(person::v::INFECTION_STATE_CHANGE_DAY, today);
+
+    // Compute how long the next state will last
+    float mean = FLAMEGPU->environment.template getProperty<float>("mean_time_to_recovered");
+    float sd = FLAMEGPU->environment.template getProperty<float>("sd_time_to_recovered");
+    float stateDuration = (FLAMEGPU->random.template normal<float>() * sd) + mean;
+    FLAMEGPU->template setVariable<float>(person::v::INFECTION_STATE_DURATION, stateDuration);
+
+    // Atomically update the number of infected individuals for the current individual's demographics when they transition into the infection state
+    totalInfectedPerDemographic[demographic_idx]++;
+
+    // Update the agent variable tracking when the "current" infection moved from exposed to infected
+    FLAMEGPU->template setVariable<std::uint32_t>(person::v::TIME_INFECTED, today);
+}
+
+/**
+ * Function for updating agent data when moved from the Infected to Recovered state
+ *
+ * Templated for due to the templated DeviceAPI object
+ *
+ * @param FLAMEGPU flamegpu device api object
+ * @param infectionStatus current infection status for the individual to be mutated in-place
+ */
+template<typename MsgIn, typename MsgOut>
+FLAMEGPU_DEVICE_FUNCTION void infectedToRecovered(flamegpu::DeviceAPI<MsgIn, MsgOut>* FLAMEGPU, disease::SEIR::InfectionStateUnderlyingType& infectionStatus) {
+    std::uint32_t today = FLAMEGPU->getStepCounter();
+
+    // Update the state
+    infectionStatus = disease::SEIR::InfectionState::Recovered;
+    FLAMEGPU->template setVariable<disease::SEIR::InfectionStateUnderlyingType>(person::v::INFECTION_STATE, infectionStatus);
+
+    // Update the day
+    FLAMEGPU->template setVariable<std::uint32_t>(person::v::INFECTION_STATE_CHANGE_DAY, today);
+
+    // Compute how long the next state will last
+    float mean = FLAMEGPU->environment.template getProperty<float>("mean_time_to_susceptible");
+    float sd = FLAMEGPU->environment.template getProperty<float>("sd_time_to_susceptible");
+    float stateDuration = (FLAMEGPU->random.template normal<float>() * sd) + mean;
+    FLAMEGPU->template setVariable<float>(person::v::INFECTION_STATE_DURATION, stateDuration);
+
+    // Update the agent variable tracking when the "current" infection moved from infected to recovered
+    FLAMEGPU->template setVariable<std::uint32_t>(person::v::TIME_RECOVERED, today);
+}
+
+/**
+ * Function for updating agent data when moved from the Recovered to Susceptible state
+ *
+ * Templated for due to the templated DeviceAPI object
+ *
+ * @param FLAMEGPU flamegpu device api object
+ * @param infectionStatus current infection status for the individual to be mutated in-place
+
+ */
+template<typename MsgIn, typename MsgOut>
+FLAMEGPU_DEVICE_FUNCTION void recoveredToSusceptible(flamegpu::DeviceAPI<MsgIn, MsgOut>* FLAMEGPU, disease::SEIR::InfectionStateUnderlyingType& infectionStatus) {
+    std::uint32_t today = FLAMEGPU->getStepCounter();
+
+    // Update the state
+    infectionStatus = disease::SEIR::Susceptible;
+    FLAMEGPU->template setVariable<disease::SEIR::InfectionStateUnderlyingType>(person::v::INFECTION_STATE, infectionStatus);
+
+    // Update the day
+    FLAMEGPU->template setVariable<std::uint32_t>(person::v::INFECTION_STATE_CHANGE_DAY, today);
+    float stateDuration = 0;  // susceptible doesn't have a fixed duration
+    FLAMEGPU->template setVariable<float>(person::v::INFECTION_STATE_DURATION, stateDuration);
+
+    // Update the agent variable tracking when the "current" infection moved from recovered to Susceptible
+    FLAMEGPU->template setVariable<std::uint32_t>(person::v::TIME_SUSCEPTIBLE, today);
+}
+
+/**
+ * Function for resetting "current" infection timing data, for use after it has been recorded
+ *
+ * Templated for due to the templated DeviceAPI object
+ *
+ * @param FLAMEGPU flamegpu device api object
+ */
+template<typename MsgIn, typename MsgOut>
+FLAMEGPU_DEVICE_FUNCTION void resetSEIRStateTimes(flamegpu::DeviceAPI<MsgIn, MsgOut>* FLAMEGPU) {
+    // Update the agent variable tracking when the "current" infection moved from Recovered to Susceptible, and the data has been logged.
+    const std::uint32_t invalidTime = static_cast<std::uint32_t>(-1);
+    // Don't reset susceptible
+    FLAMEGPU->template setVariable<std::uint32_t>(person::v::TIME_EXPOSED, invalidTime);
+    FLAMEGPU->template setVariable<std::uint32_t>(person::v::TIME_INFECTED, invalidTime);
+    FLAMEGPU->template setVariable<std::uint32_t>(person::v::TIME_RECOVERED, invalidTime);
 }
 
 }  // namespace SEIR

src/exateppabm/exatepp_abm.cu

@@ -95,7 +95,7 @@ int entrypoint(int argc, char* argv[]) {
     exateppabm::population::define(model, *config, cli_params->verbosity > 0);
 
     // Add init, step and exit functions related to data collection and output. This may want refactoring when multiple output files are supported or collected data becomes more complex.
-    exateppabm::output::define(model, cli_params->outputDir, cli_params->individualFile);
+    exateppabm::output::define(model, cli_params->outputDir, cli_params->individualFile, cli_params->transmissionFile);
 
     // Build the model control flow. This will want abstracting more in the future @todo
     // @note - not using the DAG control flow due to bugs encountered in another project when splitting between compilation units.