From 6b63009da36cc05f254298aa2c0209481310aa13 Mon Sep 17 00:00:00 2001
From: Luis Aleixo <luis.aleixo@cern.ch>
Date: Thu, 25 Jan 2024 16:38:19 +0100
Subject: [PATCH] Initial steps towards a dynamic exhalation rate

---
 caimira/models.py             | 34 ++++++++++++++++++++++++++--------
 caimira/monte_carlo/models.py |  4 ++++
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/caimira/models.py b/caimira/models.py
index 4963b80f..20db5e4e 100644
--- a/caimira/models.py
+++ b/caimira/models.py
@@ -805,6 +805,13 @@ class Activity:
 }
 
 
+@dataclass(frozen=True)
+class ActivityPiecewiseConstant(PiecewiseConstant):
+
+    #: values of the function between transitions
+    values: typing.Tuple[Activity, ...]
+
+
 @dataclass(frozen=True)
 class SimplePopulation:
     """
@@ -819,7 +826,7 @@ class SimplePopulation:
     presence: typing.Union[None, Interval]
 
     #: The physical activity being carried out by the people.
-    activity: Activity
+    activity: typing.Union[Activity, ActivityPiecewiseConstant]
 
     def __post_init__(self):
         if isinstance(self.number, int):
@@ -828,6 +835,11 @@ def __post_init__(self):
         else:
             if self.presence is not None:
                 raise TypeError(f'The presence argument must be None for a IntPiecewiseConstant number')
+            
+        if isinstance(self.activity, Activity):
+            if not isinstance(self.presence, Interval):
+                raise TypeError(f'The presence argument must be an "Interval". Got {type(self.presence)}')
+
 
     def presence_interval(self):
         if isinstance(self.presence, Interval):
@@ -1054,7 +1066,7 @@ def min_background_concentration(self) -> _VectorisedFloat:
         """
         return self.data_registry.concentration_model['min_background_concentration'] # type: ignore
 
-    def normalization_factor(self) -> _VectorisedFloat:
+    def normalization_factor(self, time: typing.Optional[float] = None) -> _VectorisedFloat:
         """
         Normalization factor (in the same unit as the concentration).
         This factor is applied to the normalized concentration only
@@ -1084,7 +1096,7 @@ def _normed_concentration_limit(self, time: float) -> _VectorisedFloat:
             invRR = np.nan if RR == 0. else 1. / RR # type: ignore
 
         return (self.population.people_present(time) * invRR / V +
-                self.min_background_concentration()/self.normalization_factor())
+                self.min_background_concentration()/self.normalization_factor(time))
 
     @method_cache
     def state_change_times(self) -> typing.List[float]:
@@ -1155,7 +1167,7 @@ def _normed_concentration(self, time: float) -> _VectorisedFloat:
         # The model always starts at t=0, but we avoid running concentration calculations
         # before the first presence as an optimisation.
         if time <= self._first_presence_time():
-            return self.min_background_concentration()/self.normalization_factor()
+            return self.min_background_concentration()/self.normalization_factor(time)
 
         next_state_change_time = self._next_state_change(time)
         RR = self.removal_rate(next_state_change_time)
@@ -1187,7 +1199,7 @@ def concentration(self, time: float) -> _VectorisedFloat:
         to this method.
         """
         return (self._normed_concentration_cached(time) *
-                self.normalization_factor())
+                self.normalization_factor(time))
 
     @method_cache
     def normed_integrated_concentration(self, start: float, stop: float) -> _VectorisedFloat:
@@ -1252,7 +1264,7 @@ def population(self) -> InfectedPopulation:
     def virus(self) -> Virus:
         return self.infected.virus
 
-    def normalization_factor(self) -> _VectorisedFloat:
+    def normalization_factor(self, time: typing.Optional[int] = None) -> _VectorisedFloat:
         # we normalize by the emission rate
         return self.infected.emission_rate_per_person_when_present()
 
@@ -1304,10 +1316,16 @@ def min_background_concentration(self) -> _VectorisedFloat:
         """
         return self.CO2_atmosphere_concentration
 
-    def normalization_factor(self) -> _VectorisedFloat:
+    def normalization_factor(self, time: typing.Optional[int] = None) -> _VectorisedFloat:
         # normalization by the CO2 exhaled per person.
         # CO2 concentration given in ppm, hence the 1e6 factor.
-        return (1e6*self.population.activity.exhalation_rate
+        if isinstance(self.population.activity, ActivityPiecewiseConstant):
+            activity: Activity = self.population.activity.value(time)
+            exhalation_rate = activity.exhalation_rate
+        else:
+            exhalation_rate = self.population.activity.exhalation_rate
+
+        return (1e6*exhalation_rate
                 *self.CO2_fraction_exhaled)
 
 
diff --git a/caimira/monte_carlo/models.py b/caimira/monte_carlo/models.py
index 7215db16..5af46593 100644
--- a/caimira/monte_carlo/models.py
+++ b/caimira/monte_carlo/models.py
@@ -81,6 +81,10 @@ def _build_mc_model(model: dataclass_instance) -> typing.Type[MCModelBase[_Model
             elif new_field.type == typing.Union[caimira.models.Interval, None]:
                 I = getattr(sys.modules[__name__], "Interval")
                 field_type = typing.Union[None, caimira.models.Interval, I]
+            
+            elif new_field.type == typing.Union[caimira.models.Activity, caimira.models.ActivityPiecewiseConstant]:
+                APC = getattr(sys.modules[__name__], "ActivityPiecewiseConstant")
+                field_type = typing.Union[caimira.models.Activity, caimira.models.ActivityPiecewiseConstant, APC]
 
             else:
                 # Check that we don't need to do anything with this type.