From b8422aaf1bf805697e0a7abcf2e671a73f54264d Mon Sep 17 00:00:00 2001
From: Nicolas Mounet <nicolas.mounet@cern.ch>
Date: Thu, 27 May 2021 13:40:46 +0200
Subject: [PATCH] Introducing _ExpirationBase and MultipleExpiration classes;
 adapting tests and model_generator accordingly (removing now obsolete
 expiration_blend function)

---
 cara/apps/calculator/model_generator.py | 32 +++------------
 cara/models.py                          | 53 +++++++++++++++++++++----
 2 files changed, 51 insertions(+), 34 deletions(-)

diff --git a/cara/apps/calculator/model_generator.py b/cara/apps/calculator/model_generator.py
index fb34392d..ad65795e 100644
--- a/cara/apps/calculator/model_generator.py
+++ b/cara/apps/calculator/model_generator.py
@@ -537,38 +537,16 @@ class FormData:
         )
 
 
-def build_expiration(expiration_definition) -> models.Expiration:
+def build_expiration(expiration_definition) -> models._ExpirationBase:
     if isinstance(expiration_definition, str):
-        return models.Expiration.types[expiration_definition]
+        return models._ExpirationBase.types[expiration_definition]
     elif isinstance(expiration_definition, dict):
-        return expiration_blend({
-            build_expiration(exp): amount
-            for exp, amount in expiration_definition.items()
-            }
+        return models.MultipleExpiration(
+            tuple([build_expiration(exp) for exp in expiration_definition.keys()]),
+            tuple(expiration_definition.values())
         )
 
 
-def expiration_blend(expiration_weights: typing.Dict[models.Expiration, int]) -> models.Expiration:
-    """
-    Combine together multiple types of Expiration, using a weighted mean to
-    compute their ejection factor and particle sizes.
-
-    """
-    ejection_factor = np.zeros(4)
-    particle_sizes = np.zeros(4)
-
-    total_weight = 0
-    for expiration, weight in expiration_weights.items():
-        total_weight += weight
-        ejection_factor += np.array(expiration.ejection_factor) * weight
-        particle_sizes += np.array(expiration.particle_sizes) * weight
-
-    r_ejection_factor: typing.Tuple[float, float, float, float] = tuple(ejection_factor/total_weight)  # type: ignore
-    r_particle_sizes: typing.Tuple[float, float, float, float] = tuple(particle_sizes/total_weight)  # type: ignore
-
-    return models.Expiration(ejection_factor=r_ejection_factor, particle_sizes=r_particle_sizes)
-
-
 def model_from_form(form: FormData) -> models.ExposureModel:
     # Initializes room with volume either given directly or as product of area and height
     if form.volume_type == 'room_volume_explicit':
diff --git a/cara/models.py b/cara/models.py
index 0943bd87..ecbccad3 100644
--- a/cara/models.py
+++ b/cara/models.py
@@ -511,14 +511,29 @@ Mask.types = {
 
 
 @dataclass(frozen=True)
-class Expiration:
+class _ExpirationBase:
+    """
+    Represents the expiration of aerosols by a person.
+    Subclasses of _ExpirationBase represent different models.
+    """
+    #: Pre-populated examples of Masks.
+    types: typing.ClassVar[typing.Dict[str, "_ExpirationBase"]]
+
+    def aerosols(self, mask: _MaskBase):
+        # total volume of aerosols expired (cm^3).
+        raise NotImplementedError("Subclass must implement")
+
+
+@dataclass(frozen=True)
+class Expiration(_ExpirationBase):
+    """
+    Simple model based on four different sizes of particles emitted,
+    with different ejection factors.
+    """
     ejection_factor: typing.Tuple[float, ...]
     particle_sizes: typing.Tuple[float, ...] = (0.8e-4, 1.8e-4, 3.5e-4, 5.5e-4)  # In cm.
 
-    #: Pre-populated examples of Expiration.
-    types: typing.ClassVar[typing.Dict[str, "Expiration"]]
-
-    def aerosols(self, mask: Mask):
+    def aerosols(self, mask: _MaskBase):
         def volume(diameter):
             return (4 * np.pi * (diameter/2)**3) / 3
         total = 0
@@ -529,7 +544,31 @@ class Expiration:
         return total
 
 
-Expiration.types = {
+@dataclass(frozen=True)
+class MultipleExpiration(_ExpirationBase):
+    """
+    Represents an expiration of aerosols.
+    Group together different modes of expiration, that represent
+    each the main expiration mode for a certain fraction of time (given by
+    the weights).
+
+    """
+    expirations: typing.Tuple[_ExpirationBase, ...]
+    weights: typing.Tuple[float, ...]
+
+    def __post_init__(self):
+        if len(self.expirations) != len(self.weights):
+            raise ValueError("expirations and weigths should contain the"
+                             "same number of elements")
+
+    def aerosols(self, mask: _MaskBase):
+        return np.array([
+            weight * expiration.aerosols(mask) / sum(self.weights)
+            for weight,expiration in zip(self.weights,self.expirations)
+        ]).sum(axis=0)
+
+
+_ExpirationBase.types = {
     'Breathing': Expiration((0.084, 0.009, 0.003, 0.002)),
     'Whispering': Expiration((0.11, 0.014, 0.004, 0.002)),
     'Talking': Expiration((0.236, 0.068, 0.007, 0.011)),
@@ -585,7 +624,7 @@ class InfectedPopulation(Population):
     virus: Virus
 
     #: The type of expiration that is being emitted whilst doing the activity.
-    expiration: Expiration
+    expiration: _ExpirationBase
 
     def emission_rate_when_present(self) -> _VectorisedFloat:
         """